From 866a069777e422b704be620e05a53349d8e789e0 Mon Sep 17 00:00:00 2001
From: Nicholas Pfaff <53228351+nepfaff@users.noreply.github.com>
Date: Tue, 1 Oct 2024 15:39:26 -0400
Subject: [PATCH] Keep depth_image on CPU (#3449)

Co-authored-by: J.Y. <132313008+jb-ye@users.noreply.github.com>
---
 nerfstudio/data/datasets/depth_dataset.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nerfstudio/data/datasets/depth_dataset.py b/nerfstudio/data/datasets/depth_dataset.py
index f2f1ab86a..d4e22e3b6 100644
--- a/nerfstudio/data/datasets/depth_dataset.py
+++ b/nerfstudio/data/datasets/depth_dataset.py
@@ -41,6 +41,8 @@ class DepthDataset(InputDataset):
         scale_factor: The scaling factor for the dataparser outputs.
     """
 
+    exclude_batch_keys_from_device = InputDataset.exclude_batch_keys_from_device + ["depth_image"]
+
     def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0):
         super().__init__(dataparser_outputs, scale_factor)
         # if there are no depth images than we want to generate them all with zoe depth