From 078dc6aa04a0484de721226a061ec74f7fce3dd7 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Sat, 23 Apr 2022 15:26:53 +0000 Subject: [PATCH] removing dotdict --- mattport/nerf/trainer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 4a1c9c5c1..e32c569af 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -58,6 +58,7 @@ class Trainer: ), num_workers=self.config.dataloader.num_workers, shuffle=True, + pin_memory=True, ) # TODO(ethan): implement the test data @@ -149,9 +150,9 @@ class Trainer: def train_iteration(self, batch: dict, step: int): """Run one iteration with a batch of inputs.""" # move batch to correct device - ray_indices = batch.indices.to(f"cuda:{self.local_rank}") + ray_indices = batch["indices"].to(f"cuda:{self.local_rank}") graph_outputs = self.graph(ray_indices) - batch.pixels = batch.pixels.to(f"cuda:{self.local_rank}") + batch["pixels"] = batch["pixels"].to(f"cuda:{self.local_rank}") losses = self.graph.get_losses(batch, graph_outputs) # or self.graph.module loss_sum, loss_dict = self.get_aggregated_loss(losses) self.optimizers.zero_grad_all()