class VideoDataset(Dataset):
    def __init__(self, video_path):
        super().__init__()
        self.video_path = video_path
        self.video_reader = imageio.get_reader(video_path)
        self.metadata = self.video_reader.get_meta_data()
        self.nframes = self.metadata["nframes"]

    def __getitem__(self, idx):
        video = imageio.get_reader(self.video_path)
        img = video.get_data(idx)
        img = torch.from_numpy(img)

        return img

    def __len__(self):
        return self.nframes

then it works