diff --git a/docs_src/examples/use_cases/example_resnet_all_in_one.py b/docs_src/examples/use_cases/example_resnet_all_in_one.py index 8c7df6fc..5d7bd264 100644 --- a/docs_src/examples/use_cases/example_resnet_all_in_one.py +++ b/docs_src/examples/use_cases/example_resnet_all_in_one.py @@ -140,7 +140,7 @@ def forward(self, x): print("Do the individual gradients match?") for name, parameter in model.named_parameters(): - match = allclose(parameter.grad_batch[sample_to_check], parameter.grad, atol=5e-6) + match = allclose(parameter.grad_batch[sample_to_check], parameter.grad, atol=1e-6) print(f"{name:>20}: {match}") if not match: raise AssertionError("Individual gradients don't match!")