Update video_dataset.py

Author: RaivoKoot

video_dataset.py

@@ -242,7 +242,7 @@ def _get(self, record, indices):
                 seg_img = self._load_image(record.path, frame_index)
                 images.extend(seg_img)
                 image_indices.append(frame_index)
-                if frame_index < record.num_frames:
+                if frame_index < record.end_frame:
                     frame_index += 1
 
         # sort images by index in case of edge cases where segments overlap each other because the overall
@@ -274,4 +274,4 @@ def forward(self, img_list):
         Returns:
             tensor of size ``NUM_IMAGES x CHANNELS x HEIGHT x WIDTH``
         """
-        return torch.stack([transforms.functional.to_tensor(pic) for pic in img_list])
+        return torch.stack([transforms.functional.to_tensor(pic) for pic in img_list])