[fbsync] update usages of torch.testing internals (#7203)

Reviewed By: vmoens

Differential Revision: D44416266

fbshipit-source-id: b37347ea1eccdca20da75483a083dff0c1fed0ae

Author: NicolasHug

test/prototype_common_utils.py

@@ -14,7 +14,7 @@
 import torch.testing
 from datasets_utils import combinations_grid
 from torch.nn.functional import one_hot
-from torch.testing._comparison import assert_equal as _assert_equal, BooleanPair, NonePair, NumberPair, TensorLikePair
+from torch.testing._comparison import BooleanPair, NonePair, not_close_error_metas, NumberPair, TensorLikePair
 from torchvision.prototype import datapoints
 from torchvision.prototype.transforms.functional import convert_dtype_image_tensor, to_image_tensor
 from torchvision.transforms.functional_tensor import _max_value as get_max_value
@@ -73,7 +73,7 @@ def compare(self) -> None:
             actual, expected = self._promote_for_comparison(actual, expected)
             mae = float(torch.abs(actual - expected).float().mean())
             if mae > self.atol:
-                raise self._make_error_meta(
+                self._fail(
                     AssertionError,
                     f"The MAE of the images is {mae}, but only {self.atol} is allowed.",
                 )
@@ -99,7 +99,7 @@ def assert_close(
     """Superset of :func:`torch.testing.assert_close` with support for PIL vs. tensor image comparison"""
     __tracebackhide__ = True
 
-    _assert_equal(
+    error_metas = not_close_error_metas(
         actual,
         expected,
         pair_types=(
@@ -117,10 +117,12 @@ def assert_close(
         check_dtype=check_dtype,
         check_layout=check_layout,
         check_stride=check_stride,
-        msg=msg,
         **kwargs,
     )
 
+    if error_metas:
+        raise error_metas[0].to_error(msg)
+
 
 assert_equal = functools.partial(assert_close, rtol=0, atol=0)
 

test/test_prototype_datasets_builtin.py

@@ -1,4 +1,3 @@
-import functools
 import io
 import pickle
 from collections import deque
@@ -9,7 +8,7 @@
 
 import torchvision.prototype.transforms.utils
 from builtin_dataset_mocks import DATASET_MOCKS, parametrize_dataset_mocks
-from torch.testing._comparison import assert_equal, ObjectPair, TensorLikePair
+from torch.testing._comparison import not_close_error_metas, ObjectPair, TensorLikePair
 
 # TODO: replace with torchdata.dataloader2.DataLoader2 as soon as it is stable-ish
 from torch.utils.data import DataLoader
@@ -25,9 +24,12 @@
 from torchvision.prototype.datasets.utils._internal import INFINITE_BUFFER_SIZE
 
 
-assert_samples_equal = functools.partial(
-    assert_equal, pair_types=(TensorLikePair, ObjectPair), rtol=0, atol=0, equal_nan=True
-)
+def assert_samples_equal(*args, msg=None, **kwargs):
+    error_metas = not_close_error_metas(
+        *args, pair_types=(TensorLikePair, ObjectPair), rtol=0, atol=0, equal_nan=True, **kwargs
+    )
+    if error_metas:
+        raise error_metas[0].to_error(msg)
 
 
 def extract_datapipes(dp):