Increase tolerance

This commit is contained in:
Simon Alibert 2024-05-02 14:45:06 +02:00
parent a66031f1f7
commit c3e57c2ae7
1 changed files with 3 additions and 3 deletions

View File

@ -255,10 +255,10 @@ def test_backward_compatibility(env_name, policy_name, extra_overrides):
output_dict, grad_stats, param_stats, actions = get_policy_stats(env_name, policy_name, extra_overrides)
for key in saved_output_dict:
assert torch.isclose(output_dict[key], saved_output_dict[key]).all()
assert torch.isclose(output_dict[key], saved_output_dict[key], rtol=1e-4, atol=1e-7).all()
for key in saved_grad_stats:
assert torch.isclose(grad_stats[key], saved_grad_stats[key]).all()
assert torch.isclose(grad_stats[key], saved_grad_stats[key], rtol=1e-4, atol=1e-7).all()
for key in saved_param_stats:
assert torch.isclose(param_stats[key], saved_param_stats[key]).all()
assert torch.isclose(param_stats[key], saved_param_stats[key], rtol=1e-4, atol=1e-7).all()
for key in saved_actions:
assert torch.isclose(actions[key], saved_actions[key], rtol=1e-4, atol=1e-7).all()