From 0eb0ceb3c07815806a9b2e51662f7ffcad031a99 Mon Sep 17 00:00:00 2001
From: Nicholas Pfaff <53228351+nepfaff@users.noreply.github.com>
Date: Tue, 25 Apr 2023 21:25:16 -0400
Subject: [PATCH] Revert loading entire dataset to GPU if using masks (#1821)

---
 nerfstudio/data/utils/dataloaders.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py
index 656202c1d..23d712d7e 100644
--- a/nerfstudio/data/utils/dataloaders.py
+++ b/nerfstudio/data/utils/dataloaders.py
@@ -116,8 +116,7 @@ class CacheDataloader(DataLoader):
         """Returns a collated batch."""
         batch_list = self._get_batch_list()
         collated_batch = self.collate_fn(batch_list)
-        exclude = [] if "mask" in collated_batch.keys() else ["image"]  # type: ignore
-        collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=exclude)
+        collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=["image"])
         return collated_batch
 
     def __iter__(self):