Update per_sample_grads.py

Author: svekars

intermediate_source/per_sample_grads.py

@@ -181,7 +181,7 @@ def compute_loss(params, buffers, sample, target):
             ft_per_sample_grad = ft_per_sample_grad.view(per_sample_grad.shape)
 
     # Use a higher tolerance for comparison
-    assert torch.allclose(per_sample_grad, ft_per_sample_grad, atol=1e-2, rtol=1e-2), \
+    assert torch.allclose(per_sample_grad, ft_per_sample_grad, atol=2e-2, rtol=2e-2), \
         f"Mismatch in {name}: max diff {(per_sample_grad - ft_per_sample_grad).abs().max().item()}"