diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py
index acc868648..3f3953558 100644
--- a/nerfstudio/fields/nerfacto_field.py
+++ b/nerfstudio/fields/nerfacto_field.py
@@ -381,8 +381,8 @@ class TorchNerfactoField(Field):
             positions = self.spatial_distortion(positions)
         else:
             positions = ray_samples.frustums.get_positions()
-        encoded_xyz = self.position_encoding(positions)
-        base_mlp_out = self.mlp_base(encoded_xyz)
+        encoded_xyz = self.position_encoding(positions.view(-1, 3)).to(torch.float32)
+        base_mlp_out = self.mlp_base(encoded_xyz).view(*ray_samples.frustums.shape, -1)
         density = self.field_output_density(base_mlp_out)
         return density, base_mlp_out
 
@@ -404,18 +404,20 @@ class TorchNerfactoField(Field):
 
         outputs = {}
         for field_head in self.field_heads:
-            encoded_dir = self.direction_encoding(ray_samples.frustums.directions)
+            directions = shift_directions_for_tcnn(ray_samples.frustums.directions)
+            directions_flat = directions.view(-1, 3)
+            encoded_dir = self.direction_encoding(directions_flat)
             mlp_out = self.mlp_head(
                 torch.cat(
                     [
-                        encoded_dir,
-                        density_embedding,  # type:ignore
+                        encoded_dir.view(-1, self.direction_encoding.get_out_dim()),
+                        density_embedding.view(-1, self.mlp_base.get_out_dim()),  # type:ignore
                         embedded_appearance.view(-1, self.appearance_embedding_dim),
                     ],
                     dim=-1,  # type:ignore
                 )
             )
-            outputs[field_head.field_head_name] = field_head(mlp_out)
+            outputs[field_head.field_head_name] = field_head(mlp_out).view(*outputs_shape, -1).to(directions)
         return outputs