diff --git a/nerfstudio/fields/vanilla_nerf_field.py b/nerfstudio/fields/vanilla_nerf_field.py index c9d2bbb78..22e11a25b 100644 --- a/nerfstudio/fields/vanilla_nerf_field.py +++ b/nerfstudio/fields/vanilla_nerf_field.py @@ -74,15 +74,15 @@ class NeRFField(Field): skip_connections=skip_connections, out_activation=nn.ReLU(), ) - - self.mlp_head = MLP( - in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim(), - num_layers=head_mlp_num_layers, - layer_width=head_mlp_layer_width, - out_activation=nn.ReLU(), - ) - self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) + + if field_heads: + self.mlp_head = MLP( + in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim(), + num_layers=head_mlp_num_layers, + layer_width=head_mlp_layer_width, + out_activation=nn.ReLU(), + ) self.field_heads = nn.ModuleList([field_head() for field_head in field_heads] if field_heads else []) # type: ignore for field_head in self.field_heads: field_head.set_in_dim(self.mlp_head.get_out_dim()) # type: ignore