diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index a273906ea..cf2fa468c 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -56,7 +56,7 @@ class Trainer: self.device = "cpu" if self.world_size == 0 else f"cuda:{self.local_rank}" @profiler.time_function - def setup(self, test_mode = False): + def setup(self, test_mode=False): """Setup the Trainer by calling other setup functions.""" dataset_inputs_dict = get_dataset_inputs_dict(**self.config.data.dataset) self.setup_datasets(dataset_inputs_dict, test_mode=test_mode) @@ -71,10 +71,10 @@ class Trainer: def load_test_dataset(self, dataset_inputs_dict, mode): """Helper method to load test or val dataset based on test/train mode""" self.val_image_dataset = ImageDataset( - image_filenames=dataset_inputs_dict[mode].image_filenames, - downscale_factor=dataset_inputs_dict[mode].downscale_factor, - alpha_color=dataset_inputs_dict[mode].alpha_color, - ) + image_filenames=dataset_inputs_dict[mode].image_filenames, + downscale_factor=dataset_inputs_dict[mode].downscale_factor, + alpha_color=dataset_inputs_dict[mode].alpha_color, + ) self.val_image_intrinsics = dataset_inputs_dict[mode].intrinsics self.val_image_camera_to_world = dataset_inputs_dict[mode].camera_to_world @@ -255,4 +255,4 @@ class Trainer: image = self.val_image_dataset[image_idx]["image"].to(self.device) psnr = self.graph.log_test_image_outputs(image_idx, step, image, outputs) self.graph.train() - return psnr \ No newline at end of file + return psnr