diff --git a/data/dataset.py b/data/dataset.py
index 0c1f8b0..5da5176 100644
--- a/data/dataset.py
+++ b/data/dataset.py
@@ -85,7 +85,7 @@ class Dataset(Dataset):
 
         start = time.perf_counter()
         # TODO: is it sure I need this K, and not some transformed one?
-        normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K)
+        normal_image = get_normals_from_depth(depth_image[None,None,:,:], depth_mask[None,:,:], K)
         print(f'normals time: {time.perf_counter() - start}')
 
         depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask
diff --git a/utils/data_utils.py b/utils/data_utils.py
index 1ef5999..d28e64f 100644
--- a/utils/data_utils.py
+++ b/utils/data_utils.py
@@ -46,10 +46,11 @@ def read_frames_file(file):
 
 
 def get_normals_from_depth(depth_image, depth_mask, camera_intrinsic):
-    depth_dilated = kornia.morphology.dilation(depth_image[:,None,:,:], torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
+    depth_dilated = kornia.morphology.dilation(depth_image, torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
     depth_dilated = kornia.morphology.dilation(depth_dilated, torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
     depth_dilated = torch.permute(depth_dilated, [0,2,3,1])
-    depth_dilated[depth_mask,:] = depth_image[:,:,:,None][depth_mask,:]
+    depth_image = torch.permute(depth_image, [0,2,3,1])
+    depth_dilated[depth_mask,:] = depth_image[depth_mask,:]
     depth_dilated = torch.permute(depth_dilated, [0,3,1,2])
     normals = kornia.geometry.depth.depth_to_normals(depth_dilated, camera_intrinsic, normalize_points=True)
     normals = torch.permute(normals, [0,2,3,1]) * depth_mask[:,:,:,None]