some tweaks and fixes

2024-05-18 22:27:05 +01:00
parent 9922c1ba85
commit 5d7960d4c7
3 changed files with 13 additions and 9 deletions
--- a/resnet.py
+++ b/resnet.py
@@ -175,7 +175,7 @@ class MaiRes(nn.Module):
        )

        self.avgpool = nn.AvgPool2d(kernel_size=RESNET_KERNEL_SIZE)
-        self.fc = nn.Linear(in_features=8192, out_features=1)
+        self.fc = nn.Linear(in_features=2048, out_features=1)

    def forward(self, x):
        x = self.conv(x)