diff --git a/models/convnext.py b/models/convnext.py
index b24ccaa..966926d 100644
--- a/models/convnext.py
+++ b/models/convnext.py
@@ -146,6 +146,8 @@ class ConvNeXtBase(nn.Module):
                     ).to(device)
                 )
 
+        self.act = nn.Tanh()
+
         for m in self.modules():
             if isinstance(m, (nn.Conv2d, nn.Linear)):
                 nn.init.trunc_normal_(m.weight, std=0.02)
@@ -160,7 +162,7 @@ class ConvNeXtBase(nn.Module):
             features.append(x)
             x = self.downsamples[i](x)
         x = self.stages[-1](x)
-        
+        x = self.act(x)
 
         return x, features
 
@@ -242,8 +244,11 @@ class ConvNeXtReverseBase2(nn.Module):
                     kernel_size=7,
                     padding=3,
                     ).to(device)
-
-                    # Dense?
+        # self.last = nn.Sequential(
+        #                 Permute([0, 2, 3, 1]),
+        #                 nn.Linear(in_features=block_setting[-1].input_channels, out_features=1, bias=True),
+        #                 Permute([0, 3, 1, 2]),
+        #             )
 
 
         for m in self.modules():
diff --git a/models/depth_estim.py b/models/depth_estim.py
index 281313f..beee526 100644
--- a/models/depth_estim.py
+++ b/models/depth_estim.py
@@ -30,9 +30,9 @@ class DepthEstimation(nn.Module):
             CNBlockConfig(9, None, 3),
         ]
 
-        self.convnext = ConvNeXt(block_setting)
+        self.convnext = ConvNeXt(block_setting, device=device)
         self.flatten = nn.Flatten()
-        self.decoder = ConvNeXt(block_setting_reverse, reversed=True)
+        self.decoder = ConvNeXt(block_setting_reverse, reversed=True, device=device)
 
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
diff --git a/test.py b/test.py
index 8b08a9a..7336d95 100644
--- a/test.py
+++ b/test.py
@@ -70,7 +70,7 @@ if __name__ == '__main__':
     train_accumulator = Accumulator()
     valid_accumulator = Accumulator()
 
-    model = DepthEstimation().to('cuda')
+    model = DepthEstimation(device=device)
 
     loss_fn = nn.MSELoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)