Adjust JetTest.test_scatter_add tolerance for TPU

sharadmv · Mar 31, 2021 · 5fe8712 · 5fe8712
1 parent d617acd
commit 5fe8712
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/jet_test.py b/tests/jet_test.py
@@ -407,7 +407,7 @@ def F(t):
 
       return grad(jacfwd(F))(0.)
 
-    self.check_jet(h, (0.,), ([1., 2., 3.],))
+    self.check_jet(h, (0.,), ([1., 2., 3.],), rtol=1e-3)
 
 
 if __name__ == '__main__':