diff --git a/data/dataset.py b/data/dataset.py
index c8eeb1a..2d5eeb9 100644
--- a/data/dataset.py
+++ b/data/dataset.py
@@ -69,6 +69,7 @@ class Dataset(Dataset):
         # body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True)
         body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, trans, zrot)
         # body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False)
+        body = torch.tensor(body.astype(np.float32), device=self.device)
         body_faces = quads2tris(body_faces)
 
         if self.random_view:
@@ -81,9 +82,9 @@ class Dataset(Dataset):
             mesh_posed = transform.transform_points(mesh_posed)
             body = transform.transform_points(body)
         else:
-            view = torch.eye(3)
+            view = torch.eye(3, device=self.device)
 
-        combined = torch.Tensor(np.concatenate((body, mesh_posed[mesh_mask, :]))).to(self.device)
+        combined = torch.concat((body, mesh_posed[mesh_mask, :]))
         # combined_no_trans = torch.Tensor(np.concatenate((body_no_trans, mesh_posed_no_trans[mesh_mask, :]))).to('cuda')
         combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to(self.device)
 
diff --git a/models/depth_estim.py b/models/depth_estim.py
index 29961c7..9d20610 100644
--- a/models/depth_estim.py
+++ b/models/depth_estim.py
@@ -13,6 +13,7 @@ class ViewEncoder(nn.Module):
         super().__init__()
         self.features: int = features
 
+        self.flatten = nn.Flatten()
         self.lin1 = nn.LazyLinear(features).to(device)
         self.act1 = nn.GELU().to(device)
         self.lin2 = nn.Linear(features, features * 15 * 20).to(device)
@@ -20,7 +21,8 @@ class ViewEncoder(nn.Module):
     
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
-        x: torch.Tensor = self.lin1(input)
+        x: torch.Tensor = self.flatten(input)
+        x= self.lin1(x)
         x = self.act1(x)
         x = self.lin2(x)
         x = self.act2(x)