diff --git a/tests/nn/pipe/test_pipe.py b/tests/nn/pipe/test_pipe.py
index f1a9f77d9..3edf13b24 100644
--- a/tests/nn/pipe/test_pipe.py
+++ b/tests/nn/pipe/test_pipe.py
@@ -20,6 +20,7 @@
 from collections import OrderedDict
 from copy import deepcopy
 import time
+import warnings
 
 import pytest
 import torch
@@ -123,23 +124,19 @@ def test_batch_size_indivisible():
     model = nn.Sequential(nn.Linear(1, 1))
     model = Pipe(model, balance=[1], devices=["cpu"], chunks=4)
 
-    with pytest.warns(None) as record:
+    with warnings.catch_warnings(action="error"):
+        # Indivisible batch size is legal.
         model(torch.rand(7, 1))
 
-    # Indivisible batch size is legal.
-    assert not record
-
 
 def test_batch_size_small():
     model = nn.Sequential(nn.Linear(1, 1))
     model = Pipe(model, balance=[1], devices=["cpu"], chunks=4)
 
-    with pytest.warns(None) as record:
+    with warnings.catch_warnings(action="error"):
+        # Batch size smaller than chunks is legal.
         model(torch.rand(2, 1))
 
-    # Batch size smaller than chunks is legal.
-    assert not record
-
 
 def test_checkpoint_mode():
     def count_grad_fn(grad_fn, name, visited=set()):
diff --git a/tests/nn/pipe_process/test_pipe.py b/tests/nn/pipe_process/test_pipe.py
index 19de13ac3..68320163d 100644
--- a/tests/nn/pipe_process/test_pipe.py
+++ b/tests/nn/pipe_process/test_pipe.py
@@ -21,6 +21,7 @@
 from copy import deepcopy
 import os
 import time
+import warnings
 
 import pytest
 import torch
@@ -222,12 +223,10 @@ def batch_size_indivisible(pipe_class):
     model = nn.Sequential(nn.Linear(1, 1))
     model = pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=4)
 
-    with pytest.warns(None) as record:
+    with warnings.catch_warnings(action="error"):
+        # Indivisible batch size is legal.
         model(torch.rand(7, 1))
 
-    # Indivisible batch size is legal.
-    assert not record
-
 
 @torch_spawn([1])
 @pytest.mark.parametrize("pipe_class", [AsyncPipe])
@@ -235,12 +234,10 @@ def batch_size_small(pipe_class):
     model = nn.Sequential(nn.Linear(1, 1))
     model = pipe_class(model, balance=[1], worker_map=get_worker_map(), chunks=4)
 
-    with pytest.warns(None) as record:
+    with warnings.catch_warnings(action="error"):
+        # Batch size smaller than chunks is legal.
         model(torch.rand(2, 1))
 
-    # Batch size smaller than chunks is legal.
-    assert not record
-
 
 @torch_spawn([1])
 @pytest.mark.parametrize("pipe_class", [AsyncPipe])