diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index ebda710b3..d7a221744 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -168,6 +168,7 @@ class Model(nn.Module): Args: camera_ray_bundle: ray bundle to calculate outputs over """ + input_device = camera_ray_bundle.directions.device num_rays_per_chunk = self.config.eval_num_rays_per_chunk image_height, image_width = camera_ray_bundle.origins.shape[:2] num_rays = len(camera_ray_bundle) @@ -176,12 +177,15 @@ class Model(nn.Module): start_idx = i end_idx = i + num_rays_per_chunk ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx) + # move the chunk inputs to the model device + ray_bundle = ray_bundle.to(self.device) outputs = self.forward(ray_bundle=ray_bundle) for output_name, output in outputs.items(): # type: ignore - if not torch.is_tensor(output): + if not isinstance(output, torch.Tensor): # TODO: handle lists of tensors as well continue - outputs_lists[output_name].append(output) + # move the chunk outputs from the model device back to the device of the inputs. + outputs_lists[output_name].append(output.to(input_device)) outputs = {} for output_name, outputs_list in outputs_lists.items(): outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1) # type: ignore