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