fix

Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
NVIDIA · Jan 28, 2025 · 20cea95 · 20cea95
1 parent 54e7438
commit 20cea95
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py
@@ -90,7 +90,7 @@ def assert_all_equal(l1: List[torch.Tensor], l2: List[torch.Tensor], names=None)
     if failed_tensors:
         print(failure_message)
         t1, t2 = failed_tensors[0]
-        torch.testing.assert_close(t1, t2, rtol=0, atol=0)
+        torch.testing.assert_close(t1, t2, rtol=0, atol=0, equal_nan=True)
 
 
 def generate_data(