diff --git a/data/dataset.py b/data/dataset.py
index 20308ed..674ebbc 100644
--- a/data/dataset.py
+++ b/data/dataset.py
@@ -14,9 +14,9 @@ import plotly.express as ex
 
 from torch.utils.data import Dataset
 
-from render import Render, DepthRenderer
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from data.render import Render, DepthRenderer
 from superSMPL.IO import readPC2Frame, writePC2
 from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth
 from cloth3d.read import DataReader
@@ -24,12 +24,13 @@ from cloth3d.util import quads2tris
 
 
 class Dataset(Dataset):
-    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt'):
+    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda'):
         self.folder = folder
         self.folder_rgb = folder_rgb
+        self.device = device
+
         self.transforms = tv.transforms.ToTensor()
         self.frames = read_frames_file(os.path.join(folder, frames_file))
-
         self.render = DepthRenderer()
         self.deodr_render : Render = Render(max_depth=10)
         self.body_faces = np.load('custom_data/smpl_faces.npy')
@@ -58,9 +59,9 @@ class Dataset(Dataset):
         # body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False)
         body_faces = quads2tris(body_faces)
 
-        combined = torch.Tensor(np.concatenate((body, mesh_posed[mesh_mask, :]))).to('cuda')
+        combined = torch.Tensor(np.concatenate((body, mesh_posed[mesh_mask, :]))).to(self.device)
         # 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('cuda')
+        combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to(self.device)
 
         start = time.perf_counter()
         # depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0]
@@ -184,7 +185,7 @@ class Dataset(Dataset):
         # ex.imshow(v2i_coverage).show()
 
 
-        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'))
+        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(self.device).to(torch.float32)
 
 
         # image = self.transforms(image)
diff --git a/models/depth_estim.py b/models/depth_estim.py
index 5f7954d..bbd3c3b 100644
--- a/models/depth_estim.py
+++ b/models/depth_estim.py
@@ -6,10 +6,10 @@ import torch.nn as nn
 
 
 class DepthEstimation(nn.Module):
-    def __init__(self):
+    def __init__(self, device='cuda'):
         super().__init__()
 
-        self.layer_1 = nn.Linear(in_features=3, out_features=3)  
+        self.layer_1 = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3, device=device, padding='same')  
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.layer_1(x)
diff --git a/test.py b/test.py
index aa40e95..b62a290 100644
--- a/test.py
+++ b/test.py
@@ -17,10 +17,10 @@ def train_step(batch):
     model.train() # train mode is on by default after construction
 
     # 1. Forward pass
-    y_pred = model(batch['mesh_posed'])
+    y_pred = model(batch['rgb_image'])
 
     # 2. Calculate loss
-    loss = loss_fn(y_pred, batch['mesh_posed'])
+    loss = loss_fn(y_pred, batch['depth_image'][:,None,:,:])
 
     # 3. Zero grad optimizer
     optimizer.zero_grad()
@@ -39,10 +39,10 @@ def valid_step(batch):
     model.eval() # put the model in evaluation mode for testing (inference)
     # 1. Forward pass
     with torch.inference_mode():
-        pred = model(batch['mesh_posed'])
+        pred = model(batch['rgb_image'])
     
         # 2. Calculate the loss
-        loss = loss_fn(pred, batch['mesh_posed'])
+        loss = loss_fn(pred, batch['depth_image'][:,None,:,:])
 
     return { 'valid/loss': loss.detach().cpu().numpy() }
 
@@ -59,8 +59,8 @@ if __name__ == '__main__':
 
     train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
     valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
-    train_dataloader = DataLoader(train_dataset, batch_size=4, drop_last=True, num_workers=2)
-    valid_dataloader = DataLoader(valid_dataset, batch_size=4, drop_last=True, num_workers=2)
+    train_dataloader = DataLoader(train_dataset, batch_size=2, drop_last=True, num_workers=0)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=2, drop_last=True, num_workers=0)
     train_accumulator = Accumulator()
     valid_accumulator = Accumulator()
 
@@ -86,8 +86,8 @@ if __name__ == '__main__':
             valid_loss = valid_step(batch)
             valid_accumulator.update(valid_loss)
 
-            print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
-            valid_accumulator.reset()
+        print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
+        valid_accumulator.reset()