Use reduce

This commit is contained in:
Matt Tancik committed 2022-08-01 17:19:51 -07:00
1 parent 3abc586427
commit fc1ba5e630
1 file changed
+1 -1
+1 -1
View File
@@ -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()