Change to linear layers for the end

Author: jacobbieker

pseudo_labeller/model/idam.py

@@ -137,7 +137,7 @@ def forward(self, x: torch.Tensor, pv_meta: torch.Tensor = None, output_latents:
         pv_meta = F.relu(pv_meta)
         # Scale down to 1 size
         x = F.adaptive_avg_pool2d(x, (1, 1))
-        x = torch.squeeze(x, dim=-1)
+        x = torch.squeeze(torch.squeeze(x, dim=-1), dim=-1)
         # Reshape to fit into 3DCNN
         x = torch.cat([x, pv_meta], dim=1)
         # Get pv_meta_output