From 667de0b7583795dadb2993a9cae27acce121b6ff Mon Sep 17 00:00:00 2001
From: Jake Austin <jake-austin@berkeley.edu>
Date: Sun, 17 Jul 2022 23:01:42 -0700
Subject: [PATCH] Subclassed abstract cached dataloader, for things like
 generic NeRFs where data fits in memory

---
 pyrad/data/dataloader.py | 8 ++++++++
 1 file changed, 8 insertions(+)

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."""