added activation for features

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():

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: