From 9922c1ba855aef6a5714b456682cda32b8542f93 Mon Sep 17 00:00:00 2001
From: Kenneth <kennethnym@outlook.com>
Date: Sat, 18 May 2024 19:58:44 +0100
Subject: [PATCH] fix reshape error between avgpool and fc layer

---
 resnet.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/resnet.py b/resnet.py
index 925f058..e3b4402 100644
--- 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