use pytest.mark.filterwarnings

Lightning-AI · Oct 18, 2024 · 809865f · 809865f
1 parent c544dc5
commit 809865f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/thunder/tests/test_dynamo.py b/thunder/tests/test_dynamo.py
@@ -517,6 +517,7 @@ def f(x, y):
 
 
 @requiresCUDA
+@pytest.mark.filterwarnings(r"ignore:`torch\.cpu\.amp\.autocast\((.*?)\)` is deprecated.*:FutureWarning")
 def test_checkpoint_converter():
     import torch.utils.checkpoint as checkpoint
 
@@ -552,10 +553,7 @@ def forward(self, x):
     out.backward(g)
 
     ref_g = torch.ones_like(ref_out)
-    with warnings.catch_warnings():
-        # FutureWarning: `torch.cpu.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cpu', args...)` instead.
-        warnings.filterwarnings("ignore", category=FutureWarning)
-        ref_out.backward(ref_g)
+    ref_out.backward(ref_g)
     torch.testing.assert_close(x.grad, x_ref.grad)
     torch.testing.assert_close(tuple(model.parameters()), tuple(ref_model.parameters()))
 

diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py
@@ -5263,7 +5263,7 @@ def _backward_checkpoint(
     from thunder.core.transforms import vjp
 
     result, grads = vjp(function)(args, grad_outputs, **kwargs)
-    return grads  # result
+    return grads
 
 
 #