debug prints

Lightning-AI · Sep 18, 2024 · 6c926d7 · 6c926d7
1 parent 6598ab8
commit 6c926d7
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tests/unittests/bases/test_ddp.py b/tests/unittests/bases/test_ddp.py
@@ -108,6 +108,10 @@ def _test_ddp_gather_autograd_different_shape(rank: int, worldsize: int = NUM_PR
  scalar2 = scalar2 + torch.sum(result[idx] * torch.ones_like(result[idx]))
  gradient1 = torch.autograd.grad(scalar1, [tensor], retain_graph=True)[0]
  gradient2 = torch.autograd.grad(scalar2, [tensor])[0]
+ print("shapes", gradient1.shape, gradient2.shape)
+ print("g1", gradient1)
+ print("g2", gradient2)
+ print("diff", gradient1 - gradient2)
  assert torch.allclose(gradient1, gradient2)