fix tensor dtypes, add test

Author: Daniel Flores

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -322,6 +322,11 @@ void SingleStreamDecoder::scanFileAndUpdateMetadataAndIndex() {
 void SingleStreamDecoder::readCustomFrameMappingsUpdateMetadataAndIndex(
     int streamIndex,
     FrameMappings customFrameMappings) {
+  TORCH_CHECK(
+      customFrameMappings.all_frames.dtype() == torch::kLong &&
+          customFrameMappings.is_key_frame.dtype() == torch::kBool &&
+          customFrameMappings.duration.dtype() == torch::kLong,
+      "all_frames and duration tensors must be int64 dtype, and is_key_frame tensor must be a bool dtype.");
   const torch::Tensor& all_frames =
       customFrameMappings.all_frames.to(torch::kLong);
   const torch::Tensor& is_key_frame =

src/torchcodec/decoders/_video_decoder.py

@@ -476,10 +476,12 @@ def _read_custom_frame_mappings(
         )
 
     frame_data = [
-        (float(frame[pts_key]), frame["key_frame"], float(frame[duration_key]))
+        (int(frame[pts_key]), frame["key_frame"], int(frame[duration_key]))
         for frame in input_data["frames"]
     ]
-    all_frames, is_key_frame, duration = map(torch.tensor, zip(*frame_data))
+    all_frames = torch.tensor([x[0] for x in frame_data], dtype=torch.int64)
+    is_key_frame = torch.tensor([x[1] for x in frame_data], dtype=torch.bool)
+    duration = torch.tensor([x[2] for x in frame_data], dtype=torch.int64)
     if not (len(all_frames) == len(is_key_frame) == len(duration)):
         raise ValueError("Mismatched lengths in frame index data")
     return all_frames, is_key_frame, duration

test/test_ops.py

@@ -481,27 +481,41 @@ def test_frame_pts_equality(self):
             assert pts_is_equal
 
     def test_seek_mode_custom_frame_mappings_fails(self):
-        decoder = create_from_file(
-            str(NASA_VIDEO.path), seek_mode="custom_frame_mappings"
-        )
         with pytest.raises(
             RuntimeError,
             match="Missing frame mappings when custom_frame_mappings seek mode is set.",
         ):
+            decoder = create_from_file(
+                str(NASA_VIDEO.path), seek_mode="custom_frame_mappings"
+            )
             add_video_stream(decoder, stream_index=0, custom_frame_mappings=None)
 
-        decoder = create_from_file(
-            str(NASA_VIDEO.path), seek_mode="custom_frame_mappings"
-        )
-        different_lengths = (
-            torch.tensor([1, 2, 3]),
-            torch.tensor([1, 2]),
-            torch.tensor([1, 2, 3]),
-        )
+        with pytest.raises(
+            RuntimeError,
+            match="all_frames and duration tensors must be int64 dtype, and is_key_frame tensor must be a bool dtype.",
+        ):
+            decoder = create_from_file(
+                str(NASA_VIDEO.path), seek_mode="custom_frame_mappings"
+            )
+            wrong_types = (
+                torch.tensor([1.1, 2.2, 3.3]),
+                torch.tensor([1, 2]),
+                torch.tensor([1, 2, 3]),
+            )
+            add_video_stream(decoder, stream_index=0, custom_frame_mappings=wrong_types)
+
         with pytest.raises(
             RuntimeError,
             match="all_frames, is_key_frame, and duration from custom_frame_mappings were not same size.",
         ):
+            decoder = create_from_file(
+                str(NASA_VIDEO.path), seek_mode="custom_frame_mappings"
+            )
+            different_lengths = (
+                torch.tensor([1, 2, 3]),
+                torch.tensor([False, False]),
+                torch.tensor([1, 2, 3]),
+            )
             add_video_stream(
                 decoder, stream_index=0, custom_frame_mappings=different_lengths
             )