diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index 1b0304d5b..2503e7e22 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -223,7 +223,7 @@ class NGPModel(Model): return metrics_dict def get_loss_dict(self, outputs, batch, metrics_dict=None): - image = batch["image"][..., :3].to(self.device) + image = batch["image"].to(self.device) pred_rgb, image = self.renderer_rgb.blend_background_for_loss_computation( pred_image=outputs["rgb"], pred_accumulation=outputs["accumulation"], diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index d19b0e2fb..6c149b4b6 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -308,7 +308,7 @@ class TensoRFModel(Model): def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: # Scaling metrics by coefficients to create the losses. device = outputs["rgb"].device - image = batch["image"][..., :3].to(device) + image = batch["image"].to(device) pred_image, image = self.renderer_rgb.blend_background_for_loss_computation( pred_image=outputs["rgb"], pred_accumulation=outputs["accumulation"],