more debugging

tests/unittests/_helpers/testers.py

@@ -32,14 +32,12 @@ def _assert_allclose(tm_result: Any, ref_result: Any, atol: float = 1e-8, key: O
     """Recursively assert that two results are within a certain tolerance."""
     # single output compare
     if isinstance(tm_result, Tensor):
-        print("TM", tm_result)
-        print("REF", ref_result)
         assert np.allclose(
             tm_result.detach().cpu().numpy() if isinstance(tm_result, Tensor) else tm_result,
             ref_result.detach().cpu().numpy() if isinstance(ref_result, Tensor) else ref_result,
             atol=atol,
             equal_nan=True,
-        )
+        ), f"TM: {tm_result}, REF: {ref_result}"
     # multi output compare
     elif isinstance(tm_result, Sequence):
         for pl_res, ref_res in zip(tm_result, ref_result):