Update TumorModel.py

TumorModel.py

@@ -16,7 +16,9 @@ class TumorClassification(nn.Module):
         self.pool3 = nn.MaxPool2d(2)
 
         self.flatten = nn.Flatten()
-        self.fc1 = nn.Linear(128 * 28 * 28, 512)  # 86528 = 128 * 28 * 28
+
+        # 👇 Update this line to match the saved .pth model
+        self.fc1 = nn.Linear(86528, 512)  # Fixed input size
         self.relu_fc = nn.ReLU()
         self.fc2 = nn.Linear(512, 256)
         self.relu_fc2 = nn.ReLU()
@@ -31,18 +33,3 @@ class TumorClassification(nn.Module):
         x = self.relu_fc2(self.fc2(x))
         x = self.output(x)
         return x
-
-
-class GliomaStageModel(nn.Module):
-    def __init__(self):
-        super(GliomaStageModel, self).__init__()
-        self.model = nn.Sequential(
-            nn.Linear(9, 128),
-            nn.ReLU(),
-            nn.Linear(128, 64),
-            nn.ReLU(),
-            nn.Linear(64, 4)
-        )
-
-    def forward(self, x):
-        return self.model(x)