Lower atol threshold for DT's save_policy test

takuseno · May 5, 2024 · b734d4e · b734d4e
1 parent 8f08fb3
commit b734d4e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/algos/transformer/algo_test.py b/tests/algos/transformer/algo_test.py
@@ -295,7 +295,7 @@ def save_policy_tester(
     if algo.get_action_type() == ActionSpace.DISCRETE:
         assert action == algo.predict(inpt).argmax()
     else:
-        assert np.allclose(action, algo.predict(inpt), atol=1e-4)
+        assert np.allclose(action, algo.predict(inpt), atol=1e-3)
 
     # check save_policy as ONNX
     algo.save_policy(os.path.join("test_data", "model.onnx"))
@@ -319,4 +319,4 @@ def save_policy_tester(
     if algo.get_action_type() == ActionSpace.DISCRETE:
         assert action == algo.predict(inpt).argmax()
     else:
-        assert np.allclose(action, algo.predict(inpt), atol=1e-4)
+        assert np.allclose(action, algo.predict(inpt), atol=1e-3)