trying to get basic model working

test.py

@@ -1,4 +1,6 @@
 import os
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+
 import sys
 
 import torch
@@ -7,9 +9,13 @@ import torch.optim as optim
 import torch.nn as nn
 
 from torch.utils.data import DataLoader
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
 from data.dataset import Dataset
 from utils.utils import Accumulator
 from models.depth_estim import DepthEstimation
+from models.convnext import ConvNeXt
 
 
 def train_step(batch):
@@ -20,7 +26,7 @@ def train_step(batch):
     y_pred = model(batch['rgb_image'])
 
     # 2. Calculate loss
-    loss = loss_fn(y_pred, batch['depth_image'][:,None,:,:])
+    loss = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
 
     # 3. Zero grad optimizer
     optimizer.zero_grad()
@@ -42,7 +48,7 @@ def valid_step(batch):
         pred = model(batch['rgb_image'])
     
         # 2. Calculate the loss
-        loss = loss_fn(pred, batch['depth_image'][:,None,:,:])
+        loss = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
 
     return { 'valid/loss': loss.detach().cpu().numpy() }
 
@@ -57,14 +63,14 @@ def batch_to_device(batch, device):
 if __name__ == '__main__':
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    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_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt')
+    valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt')
     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()
 
-    model = DepthEstimation()
+    model = DepthEstimation().to('cuda')
 
     loss_fn = nn.MSELoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)