diff --git a/data/dataset.py b/data/dataset.py
index 6153e18..a8f2eb7 100644
--- a/data/dataset.py
+++ b/data/dataset.py
@@ -1,6 +1,8 @@
 import os
 import sys
 
+os.environ["CUDA_VISIBLE_DEVICES"] = '1'
+
 import math
 import torch
 import torchvision as tv
@@ -18,14 +20,16 @@ class Dataset(Dataset):
         return len(self.frames)
 
     def __getitem__(self, idx):
-        img_path = os.path.join(self.folder, 'mesh_posed', self.frames[idx])
-        image = np.load(img_path)['mesh_posed']
+        path = os.path.join(self.folder, 'mesh_posed', self.frames[idx])
+        mesh_posed = np.load(path)['mesh_posed']
         # image = self.transforms(image)
-        return image
+        return { 'mesh_posed': mesh_posed }
 
 
 if __name__ == '__main__':
     from torch.utils.data import DataLoader
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
     dataset = Dataset('/home/workspace/local-data/train')
     train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, num_workers=2)
diff --git a/models/depth_estim.py b/models/depth_estim.py
new file mode 100644
index 0000000..5f7954d
--- /dev/null
+++ b/models/depth_estim.py
@@ -0,0 +1,17 @@
+import os
+import sys
+
+import torch
+import torch.nn as nn
+
+
+class DepthEstimation(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.layer_1 = nn.Linear(in_features=3, out_features=3)  
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.layer_1(x)
+        
+        return x
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..7cc64a1
--- /dev/null
+++ b/test.py
@@ -0,0 +1,95 @@
+import os
+import sys
+
+import torch
+
+import torch.optim as optim
+import torch.nn as nn
+
+from torch.utils.data import DataLoader
+from data.dataset import Dataset
+from utils.utils import Accumulator
+from models.depth_estim import DepthEstimation
+
+
+def train_step(batch):
+    ### Training
+    model.train() # train mode is on by default after construction
+
+    # 1. Forward pass
+    y_pred = model(batch['mesh_posed'])
+
+    # 2. Calculate loss
+    loss = loss_fn(y_pred, batch['mesh_posed'])
+
+    # 3. Zero grad optimizer
+    optimizer.zero_grad()
+
+    # 4. Loss backward
+    loss.backward()
+
+    # 5. Step the optimizer
+    optimizer.step()
+
+    return { 'train/loss': loss.detach().cpu().numpy() }
+
+
+def valid_step(batch):
+    ### Testing
+    model.eval() # put the model in evaluation mode for testing (inference)
+    # 1. Forward pass
+    with torch.inference_mode():
+        pred = model(batch['mesh_posed'])
+    
+        # 2. Calculate the loss
+        loss = loss_fn(pred, batch['mesh_posed'])
+
+    return { 'valid/loss': loss.detach().cpu().numpy() }
+
+
+def batch_to_device(batch, device):
+    filter = {'mesh_posed'}
+    for key in batch:
+        if key in filter:
+            batch[key] = batch[key].to(device)
+
+
+if __name__ == '__main__':
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+    train_dataset = Dataset('/home/workspace/local-data/train')
+    valid_dataset = Dataset('/home/workspace/local-data/valid')
+    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_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    model = DepthEstimation()
+
+    loss_fn = nn.MSELoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+    iteration = 0
+    while True:
+
+        for batch in train_dataloader:
+            train_loss = train_step(batch)
+            train_accumulator.update(train_loss)
+
+            if iteration % 1000 == 0:
+                print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}")
+                train_accumulator.reset()
+
+            iteration += 1
+
+        for batch in valid_dataloader:
+            valid_loss = valid_step(batch)
+            valid_accumulator.update(valid_loss)
+
+            print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
+            valid_accumulator.reset()
+
+
+
+    pass 
\ No newline at end of file
diff --git a/test/test.py b/test/test.py
deleted file mode 100644
index 93dc597..0000000
--- a/test/test.py
+++ /dev/null
@@ -1,9 +0,0 @@
-import os
-import sys
-
-import torch
-
-
-if __name__ == '__main__':
-
-    pass
\ No newline at end of file
diff --git a/utils/utils.py b/utils/utils.py
new file mode 100644
index 0000000..5e23680
--- /dev/null
+++ b/utils/utils.py
@@ -0,0 +1,25 @@
+import os
+import sys
+import numpy as np
+
+
+
+class Accumulator():
+    def __init__(self):
+        self.dict = dict()
+
+    def update(self, dict):
+        for key in dict.keys():
+            if key in self.dict.keys():
+                self.dict[key].append(dict[key])
+            else:
+                self.dict[key] = [dict[key]]
+
+    def get(self):
+        result = dict()
+        for key in self.dict.keys():
+            result[key] = np.mean(self.dict[key])
+        return result
+
+    def reset(self):
+        self.dict = dict()