Class FrameSamplerOverfit

flowmap/frame_sampler/frame_sampler_overfit.py:19–33

class FrameSamplerOverfit(FrameSampler[FrameSamplerOverfitCfg]):
    def sample(
        self,
        num_frames_in_video: int,
        device: torch.device,
    ) -> Int64[Tensor, " frame"]:
        start = self.cfg.start or 0
        num_frames = self.cfg.num_frames or num_frames_in_video
        step = self.cfg.step or 1
        return torch.arange(
            start,
            start + num_frames * step,
            step,
            device=device,
        )