From 8fe6d7c762c69cac7ad631e46cb5437459db26c0 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Wed, 28 Sep 2022 17:38:19 +0000 Subject: [PATCH] fix module registration from python list to ModuleList --- nerfactory/models/proposal.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/nerfactory/models/proposal.py b/nerfactory/models/proposal.py index 2a34dbc83..7a08faffb 100644 --- a/nerfactory/models/proposal.py +++ b/nerfactory/models/proposal.py @@ -109,12 +109,12 @@ class ProposalModel(Model): self.field = TCNNInstantNGPField(self.scene_bounds.aabb, spatial_distortion=scene_contraction) # Build the proposal network(s) + self.proposal_networks = torch.nn.ModuleList() if self.config.use_same_proposal_network: network = DensityField(self.scene_bounds.aabb, spatial_distortion=scene_contraction) - self.proposal_networks = [network] + self.proposal_networks.append(network) self.density_fns = [network.density_fn for _ in range(self.config.num_proposal_network_iterations)] else: - self.proposal_networks = [] for _ in range(self.config.num_proposal_network_iterations): network = DensityField(self.scene_bounds.aabb, spatial_distortion=scene_contraction) self.proposal_networks.append(network) @@ -146,10 +146,7 @@ class ProposalModel(Model): def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} - proposal_networks_params = [] - for network in self.proposal_networks: - proposal_networks_params += list(network.parameters()) - param_groups["proposal_networks"] = proposal_networks_params + param_groups["proposal_networks"] = list(self.proposal_networks.parameters()) param_groups["fields"] = list(self.field.parameters()) return param_groups