Improve numerical robustness of input diff hook unit test

synsense · Nov 13, 2024 · 347e225 · 347e225
1 parent 9a194ad
commit 347e225
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_hooks.py b/tests/test_hooks.py
@@ -139,4 +139,4 @@ def test_input_diff_hook():
             layer.register_forward_hook(hooks.input_diff_hook)
     model(inp)
     for idx, correct in correct_values.items():
-        assert (model[idx].hook_data["diff_output"] == correct).all()
+        assert torch.allclose(model[idx].hook_data["diff_output"], correct)