Improve io.videoreader pyav backend performance (#8173)

6640e494 · Ivan Magazinnik · GitHub · ffc58ef4 · 6640e494
Unverified Commit 6640e494 authored Dec 20, 2023 by Ivan Magazinnik Committed by GitHub Dec 20, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

torchvision/io/video_reader.py torchvision/io/video_reader.py +2 -2

No files found.
--- a/torchvision/io/video_reader.py
+++ b/torchvision/io/video_reader.py
@@ -192,9 +192,9 @@ class VideoReader:
                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: