diff --git a/nerfstudio/utils/tensor_dataclass.py b/nerfstudio/utils/tensor_dataclass.py index a2b8d1dadb..1710881adc 100644 --- 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 index 8971c83f88..10e21adc05 100644 --- 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()