[bugfix] Correct call to torch.no_grad (#5124)

Co-authored-by: Gregor Koporec <[email protected]> Co-authored-by: Jirka Borovec <[email protected]> Co-authored-by: Sean Naren <[email protected]>
Lightning-AI · Dec 19, 2020 · 4c34855 · 4c34855
1 parent 5d2fa98
commit 4c34855
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
@@ -206,6 +206,6 @@ def all_gather_ddp_if_available(
         if sync_grads:
             return AllGatherGrad.apply(tensor, group)
         else:
-            with torch.no_grad:
+            with torch.no_grad():
                 return AllGatherGrad.apply(tensor, group)
     return tensor