Fix test_write_video_with_audio (#8553)

pytorch · Jul 30, 2024 · 845bb61 · 845bb61
1 parent 120e22b
commit 845bb61
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/test_io.py b/test/test_io.py
@@ -6,7 +6,7 @@
 import pytest
 import torch
 import torchvision.io as io
-from common_utils import assert_equal
+from common_utils import assert_equal, cpu_and_cuda
 from torchvision import get_video_backend
 
 
@@ -255,7 +255,7 @@ def test_read_video_partially_corrupted_file(self):
                 assert_equal(video, data)
 
     @pytest.mark.skipif(sys.platform == "win32", reason="temporarily disabled on Windows")
-    @pytest.mark.parametrize("device", ["cpu", "cuda"])
+    @pytest.mark.parametrize("device", cpu_and_cuda())
     def test_write_video_with_audio(self, device, tmpdir):
         f_name = os.path.join(VIDEO_DIR, "R6llTwEh07w.mp4")
         video_tensor, audio_tensor, info = io.read_video(f_name, pts_unit="sec")