Improve io.videoreader pyav backend performance (#8173)

pytorch · Dec 20, 2023 · 6640e49 · 6640e49
1 parent ffc58ef
commit 6640e49
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchvision/io/video_reader.py b/torchvision/io/video_reader.py
@@ -192,9 +192,9 @@ def __next__(self) -> Dict[str, Any]:
                 frame = next(self._c)
                 pts = float(frame.pts * frame.time_base)
                 if "video" in self.pyav_stream:
-                    frame = torch.tensor(frame.to_rgb().to_ndarray()).permute(2, 0, 1)
+                    frame = torch.as_tensor(frame.to_rgb().to_ndarray()).permute(2, 0, 1)
                 elif "audio" in self.pyav_stream:
-                    frame = torch.tensor(frame.to_ndarray()).permute(1, 0)
+                    frame = torch.as_tensor(frame.to_ndarray()).permute(1, 0)
                 else:
                     frame = None
             except av.error.EOFError: