test_autocast_torch_matmul: fix flakiness (#851)

Lightning-AI · Jul 24, 2024 · 1c59321 · 1c59321
1 parent bb29e5a
commit 1c59321
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/thunder/tests/test_autocast.py b/thunder/tests/test_autocast.py
@@ -255,7 +255,7 @@ def foo(a, b):
     torch.testing.assert_close(actual, expected)
 
     if requires_grad:
-        go = torch.randn_like(expected)
+        go = torch.ones_like(expected) / expected.numel()
         eager_grads = torch.autograd.grad(expected, [a, b], go)
         jit_grads = torch.autograd.grad(actual, [a, b], go)