2025-09-09 nightly release (ee77f57)

Author: pytorchbot

src/torchcodec/decoders/_video_decoder.py

@@ -247,6 +247,12 @@ def get_frames_at(self, indices: list[int]) -> FrameBatch:
         Returns:
             FrameBatch: The frames at the given indices.
         """
+        if isinstance(indices, torch.Tensor):
+            # TODO we should avoid converting tensors to lists and just let the
+            # core ops and C++ code natively accept tensors.  See
+            # https://github.com/pytorch/torchcodec/issues/879
+            indices = indices.to(torch.int).tolist()
+
         data, pts_seconds, duration_seconds = core.get_frames_at_indices(
             self._decoder, frame_indices=indices
         )
@@ -322,6 +328,12 @@ def get_frames_played_at(self, seconds: list[float]) -> FrameBatch:
         Returns:
             FrameBatch: The frames that are played at ``seconds``.
         """
+        if isinstance(seconds, torch.Tensor):
+            # TODO we should avoid converting tensors to lists and just let the
+            # core ops and C++ code natively accept tensors.  See
+            # https://github.com/pytorch/torchcodec/issues/879
+            seconds = seconds.to(torch.float).tolist()
+
         data, pts_seconds, duration_seconds = core.get_frames_by_pts(
             self._decoder, timestamps=seconds
         )

test/resources/nasa_13013.mp4.stream3.frame000289.pt

@@ -180,7 +180,7 @@ def test_getitem_int(self, num_ffmpeg_threads, device, seek_mode):
         ref_frame0 = NASA_VIDEO.get_frame_data_by_index(0).to(device)
         ref_frame1 = NASA_VIDEO.get_frame_data_by_index(1).to(device)
         ref_frame180 = NASA_VIDEO.get_frame_data_by_index(180).to(device)
-        ref_frame_last = NASA_VIDEO.get_frame_data_by_index(289).to(device)
+        ref_frame_last = NASA_VIDEO.get_frame_data_by_index(389).to(device)
 
         assert_frames_equal(ref_frame0, decoder[0])
         assert_frames_equal(ref_frame1, decoder[1])
@@ -193,7 +193,7 @@ def test_getitem_numpy_int(self):
         ref_frame0 = NASA_VIDEO.get_frame_data_by_index(0)
         ref_frame1 = NASA_VIDEO.get_frame_data_by_index(1)
         ref_frame180 = NASA_VIDEO.get_frame_data_by_index(180)
-        ref_frame_last = NASA_VIDEO.get_frame_data_by_index(289)
+        ref_frame_last = NASA_VIDEO.get_frame_data_by_index(389)
 
         # test against numpy.int64
         assert_frames_equal(ref_frame0, decoder[numpy.int64(0)])
@@ -404,7 +404,7 @@ def test_iteration(self, device, seek_mode):
         ref_frame9 = NASA_VIDEO.get_frame_data_by_index(9).to(device)
         ref_frame35 = NASA_VIDEO.get_frame_data_by_index(35).to(device)
         ref_frame180 = NASA_VIDEO.get_frame_data_by_index(180).to(device)
-        ref_frame_last = NASA_VIDEO.get_frame_data_by_index(289).to(device)
+        ref_frame_last = NASA_VIDEO.get_frame_data_by_index(389).to(device)
 
         # Access an arbitrary frame to make sure that the later iteration
         # still works as expected. The underlying C++ decoder object is
@@ -1390,6 +1390,17 @@ def test_custom_frame_mappings_init_fails_invalid_json(self, tmp_path, device):
                         custom_frame_mappings=custom_frame_mappings,
                     )
 
+    def test_get_frames_at_tensor_indices(self):
+        # Non-regression test for tensor support in get_frames_at() and
+        # get_frames_played_at()
+        decoder = VideoDecoder(NASA_VIDEO.path)
+
+        decoder.get_frames_at(torch.tensor([0, 10], dtype=torch.int))
+        decoder.get_frames_at(torch.tensor([0, 10], dtype=torch.float))
+
+        decoder.get_frames_played_at(torch.tensor([0, 1], dtype=torch.int))
+        decoder.get_frames_played_at(torch.tensor([0, 1], dtype=torch.float))
+
 
 class TestAudioDecoder:
     @pytest.mark.parametrize("asset", (NASA_AUDIO, NASA_AUDIO_MP3, SINE_MONO_S32))

test/test_decoders.py

@@ -337,7 +337,7 @@ def test_throws_exception_at_eof(self, device):
 
         seek_to_pts(decoder, 12.979633)
         last_frame, _, _ = get_next_frame(decoder)
-        reference_last_frame = NASA_VIDEO.get_frame_data_by_index(289)
+        reference_last_frame = NASA_VIDEO.get_frame_data_by_index(389)
         assert_frames_equal(last_frame, reference_last_frame.to(device))
         with pytest.raises(IndexError, match="no more frames"):
             get_next_frame(decoder)
@@ -1059,7 +1059,7 @@ def seek(self, offset: int, whence: int) -> int:
         seek_to_pts(decoder, 12.979633)
 
         frame_last, *_ = get_next_frame(decoder)
-        reference_frame_last = NASA_VIDEO.get_frame_data_by_index(289)
+        reference_frame_last = NASA_VIDEO.get_frame_data_by_index(389)
         assert_frames_equal(frame_last, reference_frame_last.to(device))
 
         assert file_counter.num_seeks > initialization_seeks