From fc1ba5e63027e9f2a6b9d3ff0bb5a7e32cf57d15 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Mon, 1 Aug 2022 17:19:51 -0700 Subject: [PATCH] Use reduce --- nerfactory/engine/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 0662f6379..2f98a0f7a 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -212,7 +212,7 @@ class Trainer: self.optimizers.zero_grad_all() with torch.autocast(device_type=ray_indices.device.type, enabled=self.mixed_precision): _, loss_dict, metrics_dict = self.graph.forward(ray_indices=ray_indices, batch=batch) - loss = typing.cast(torch.Tensor, sum(loss_dict.values())) + loss = functools.reduce(torch.add, loss_dict.values()) self.grad_scaler.scale(loss).backward() # type: ignore self.optimizers.optimizer_scaler_step_all(self.grad_scaler) self.grad_scaler.update()