From 3e60dbd590d1aef53e443d0d2dcb792a91afe481 Mon Sep 17 00:00:00 2001 From: Brizar <1500595+bmmtstb@users.noreply.github.com> Date: Thu, 25 Jul 2024 16:03:21 +0200 Subject: [PATCH] Automatically send video to CPU in io.write_video (#8537) Co-authored-by: Nicolas Hug --- test/test_io.py | 6 +++++- torchvision/io/video.py | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/test/test_io.py b/test/test_io.py index c45180571f0..0f57da29c27 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -255,10 +255,14 @@ def test_read_video_partially_corrupted_file(self): assert_equal(video, data) @pytest.mark.skipif(sys.platform == "win32", reason="temporarily disabled on Windows") - def test_write_video_with_audio(self, tmpdir): + @pytest.mark.parametrize("device", ["cpu", "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") + video_tensor = video_tensor.to(device) + audio_tensor = audio_tensor.to(device) + out_f_name = os.path.join(tmpdir, "testing.mp4") io.video.write_video( out_f_name, diff --git a/torchvision/io/video.py b/torchvision/io/video.py index 9b2eacbab11..3bccf414a71 100644 --- a/torchvision/io/video.py +++ b/torchvision/io/video.py @@ -80,7 +80,7 @@ def write_video( if not torch.jit.is_scripting() and not torch.jit.is_tracing(): _log_api_usage_once(write_video) _check_av_available() - video_array = torch.as_tensor(video_array, dtype=torch.uint8).numpy() + video_array = torch.as_tensor(video_array, dtype=torch.uint8).numpy(force=True) # PyAV does not support floating point numbers with decimal point # and will throw OverflowException in case this is not the case