diff --git a/pyrad/data/dataloader.py b/pyrad/data/dataloader.py index 50f97171e..1b6489b10 100644 --- a/pyrad/data/dataloader.py +++ b/pyrad/data/dataloader.py @@ -20,6 +20,7 @@ from abc import abstractmethod import random from typing import Dict, List, Tuple, Union from omegaconf import ListConfig +import torch from torchtyping import TensorType @@ -191,6 +192,13 @@ class AbstractDataloaderV2(nn.Module): """Returns the next batch of data from the eval dataloader""" +class AbstractCachedDataloaderV2(AbstractDataloaderV2): + """Subclass of the new V2 dataloader that is used for when things fit in memory""" + + camera_to_world: torch.Tensor + intrinsics: torch.Tensor + + class TrainDataloader: """Training dataloader base class."""