tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-13 11:24:00 +08:00
Code Activity

Revert loading entire dataset to GPU if using masks (#1821)

Browse Source
This commit is contained in:
Nicholas Pfaff authored and GitHub committed 2023-04-25 21:25:16 -04:00
1 parent d3f76f7f3c
commit 0eb0ceb3c0
1 file changed
+1 -2
Download Patch File Download Diff File Expand all files Collapse all files
nerfstudio/data/utils/dataloaders.py
+1 -2
View File
@@ -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):