4
0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00

Small device bug in mipnerf (#506)

device bug in mipnerf.
This commit is contained in:
Alexander Kristoffersen 2022-09-21 16:35:41 -07:00 committed by GitHub
parent 7b5a251a30
commit a1152be22b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 2 additions and 2 deletions

View File

@ -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}

View File

@ -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 = {