basic test

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)