Fix allclose check in case of NaNs

pyg-team · Oct 15, 2022 · 2febad4 · 2febad4
1 parent 489d5f3
commit 2febad4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/nn/dense/test_linear.py b/test/nn/dense/test_linear.py
@@ -94,7 +94,8 @@ def test_copy_linear(lazy):
         assert torch.allclose(copied_lin.weight, lin.weight)
     assert id(copied_lin.bias) != id(lin.bias)
     assert copied_lin.bias.data_ptr() != lin.bias.data_ptr()
-    assert torch.allclose(copied_lin.bias, lin.bias, atol=1e-6)
+    if int(torch.isnan(lin.bias).sum()) == 0:
+        assert torch.allclose(copied_lin.bias, lin.bias)
 
 
 def test_hetero_linear():