tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-10 18:17:31 +08:00
Code

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

Browse Source 
* Don't throw away non-tensor objects in TensorDataclass

* Update test_tensor_dataclass.py

---------

Co-authored-by: Justin Kerr <justin.g.kerr@gmail.com>
This commit is contained in:
Brent Yi 2024-01-19 01:18:19 -08:00 committed by GitHub
parent f55bd04ca6
commit 99073eb187
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nerfstudio/utils/tensor_dataclass.py
View File

@ -141,6 +141,8 @@ class TensorDataclass:
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:

6
tests/utils/test_tensor_dataclass.py
View File

@ -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()