added activation for features

2023-01-10 13:37:31 +00:00
parent bca2e0bd2c
commit 337ff40c89
3 changed files with 11 additions and 6 deletions
@@ -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: