[Tests] Better prints (huggingface#1043)

utils/__init__.py

@@ -40,7 +40,15 @@
 
 
 if is_torch_available():
-    from .testing_utils import floats_tensor, load_image, parse_flag_from_env, require_torch_gpu, slow, torch_device
+    from .testing_utils import (
+        floats_tensor,
+        load_image,
+        parse_flag_from_env,
+        require_torch_gpu,
+        slow,
+        torch_all_close,
+        torch_device,
+    )
 
 
 logger = get_logger(__name__)

utils/testing_utils.py

@@ -34,6 +34,14 @@
         torch_device = "mps" if (mps_backend_registered and torch.backends.mps.is_available()) else torch_device
 
 
+def torch_all_close(a, b, *args, **kwargs):
+    if not is_torch_available():
+        raise ValueError("PyTorch needs to be installed to use this function.")
+    if not torch.allclose(a, b, *args, **kwargs):
+        assert False, f"Max diff is absolute {(a - b).abs().max()}. Diff tensor is {(a - b).abs()}."
+    return True
+
+
 def get_tests_dir(append_path=None):
     """
     Args: