diff --git a/nerfactory/models/mipnerf.py b/nerfactory/models/mipnerf.py
index 5a08addfe..dccd85a43 100644
--- a/nerfactory/models/mipnerf.py
+++ b/nerfactory/models/mipnerf.py
@@ -137,7 +137,7 @@ class MipNerfModel(Model):
         return outputs
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
-        image = batch["image"]
+        image = batch["image"].to(self.device)
         rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"])
         rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"])
         loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine}
diff --git a/nerfactory/models/mipnerf_360.py b/nerfactory/models/mipnerf_360.py
index 9390b1e82..119379624 100644
--- a/nerfactory/models/mipnerf_360.py
+++ b/nerfactory/models/mipnerf_360.py
@@ -147,7 +147,7 @@ class MipNerf360Model(Model):
         return outputs
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
-        image = batch["image"]
+        image = batch["image"].to(self.device)
         rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"])
         rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"])
         loss_dict = {