Make sure device is general in test

pytorch · Jan 10, 2025 · 25437b0 · 25437b0
1 parent d49fde5
commit 25437b0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/decoders/test_video_decoder.py b/test/decoders/test_video_decoder.py
@@ -424,7 +424,7 @@ def test_get_frame_at_av1(self, device):
         decoded_frame10 = decoder.get_frame_at(10)
         assert decoded_frame10.duration_seconds == ref_frame_info10.duration_seconds
         assert decoded_frame10.pts_seconds == ref_frame_info10.pts_seconds
-        assert_frames_equal(decoded_frame10.data, ref_frame10.to(device="cpu"))
+        assert_frames_equal(decoded_frame10.data, ref_frame10.to(device=device))
 
     @pytest.mark.parametrize("device", cpu_and_cuda())
     def test_get_frame_played_at(self, device):