Don't throw away non-tensor objects in TensorDataclass (#2783)

* Don't throw away non-tensor objects in TensorDataclass * Update test_tensor_dataclass.py --------- Co-authored-by: Justin Kerr <[email protected]>
nerfstudio-project · Jan 19, 2024 · 99073eb · 99073eb
1 parent f55bd04
commit 99073eb
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/nerfstudio/utils/tensor_dataclass.py b/nerfstudio/utils/tensor_dataclass.py
@@ -141,6 +141,8 @@ def _broadcast_dict_fields(self, dict_: Dict, batch_shape) -> Dict:
                 new_dict[k] = v.broadcast_to(batch_shape)
             elif isinstance(v, Dict):
                 new_dict[k] = self._broadcast_dict_fields(v, batch_shape)
+            else:
+                new_dict[k] = v
         return new_dict
 
     def __getitem__(self: TensorDataclassT, indices) -> TensorDataclassT:

diff --git a/tests/utils/test_tensor_dataclass.py b/tests/utils/test_tensor_dataclass.py
@@ -179,6 +179,12 @@ def test_iter():
         assert batch.b.shape == (4, 5)
 
 
+def test_non_tensor():
+    """Test iterating over tensor dataclass"""
+    # We shouldn't throw away non-dataclass values.
+    assert DummyTensorDataclass(a=torch.ones((3, 10)), b={"k": 2}, c=None).b == {"k": 2}  # type: ignore
+
+
 if __name__ == "__main__":
     test_init()
     test_broadcasting()