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

Delete useless lists of tensors put to CUDA which may cause OOM. (#1748)

Browse Source 
eval_large_imgs
This commit is contained in:
Yuantao Chen 2023-04-18 14:35:34 +08:00 committed by GitHub
parent e48468cfae
commit ce733a8308
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
nerfstudio/models/base_model.py
View File

@ -177,12 +177,12 @@ class Model(nn.Module):
ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx)
outputs = self.forward(ray_bundle=ray_bundle)
for output_name, output in outputs.items(): # type: ignore
if not torch.is_tensor(output):
# TODO: handle lists of tensors as well
continue
outputs_lists[output_name].append(output)
outputs = {}
for output_name, outputs_list in outputs_lists.items():
if not torch.is_tensor(outputs_list[0]):
# TODO: handle lists of tensors as well
continue
outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1) # type: ignore
return outputs