fix reshape error between avgpool and fc layer

2024-05-18 19:58:44 +01:00
parent a61e12c5db
commit 9922c1ba85
1 changed files with 1 additions and 1 deletions
--- a/resnet.py
+++ b/resnet.py
@@ -187,7 +187,7 @@ class MaiRes(nn.Module):
        x = self.layer_blue(x)

        x = self.avgpool(x)
-        x = torch.flatten(x)
+        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x