Fix test_write_video_with_audio.. again (#8558)

pytorch · Jul 30, 2024 · 09077a8 · 09077a8
1 parent 478c670
commit 09077a8
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 7 deletions.
diff --git a/test/test_io.py b/test/test_io.py
@@ -260,17 +260,14 @@ 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,
-            video_tensor,
+            video_tensor.to(device),
             round(info["video_fps"]),
             video_codec="libx264rgb",
             options={"crf": "0"},
-            audio_array=audio_tensor,
+            audio_array=audio_tensor.to(device),
             audio_fps=info["audio_fps"],
             audio_codec="aac",
         )

diff --git a/torchvision/datasets/kinetics.py b/torchvision/datasets/kinetics.py
@@ -16,7 +16,7 @@
 from .vision import VisionDataset
 
 
-def _dl_wrap(tarpath: Union[str, Path], videopath:  Union[str, Path], line: str) -> None:
+def _dl_wrap(tarpath: Union[str, Path], videopath: Union[str, Path], line: str) -> None:
     download_and_extract_archive(line, tarpath, videopath)
 
 

diff --git a/torchvision/io/video.py b/torchvision/io/video.py
@@ -115,7 +115,7 @@ def write_video(
             audio_sample_fmt = container.streams.audio[0].format.name
 
             format_dtype = np.dtype(audio_format_dtypes[audio_sample_fmt])
-            audio_array = torch.as_tensor(audio_array).numpy().astype(format_dtype)
+            audio_array = torch.as_tensor(audio_array).numpy(force=True).astype(format_dtype)
 
             frame = av.AudioFrame.from_ndarray(audio_array, format=audio_sample_fmt, layout=audio_layout)