Update TumorModel.py

TumorModel.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-# 🎯 Tumor Type Classification Model
+# 🧠 Tumor Type Classification Model
 class TumorClassification(nn.Module):
     def __init__(self):
         super(TumorClassification, self).__init__()
@@ -17,7 +17,7 @@ class TumorClassification(nn.Module):
         self.pool3 = nn.MaxPool2d(2)
 
         self.flatten = nn.Flatten()
-        self.fc1 = nn.Linear(86528, 512)  # fixed to match pth
+        self.fc1 = nn.Linear(86528, 512)  # Adjust this number to match your original
         self.relu_fc = nn.ReLU()
         self.fc2 = nn.Linear(512, 256)
         self.relu_fc2 = nn.ReLU()
@@ -32,17 +32,20 @@ class TumorClassification(nn.Module):
         x = self.relu_fc2(self.fc2(x))
         return self.output(x)
 
-# 🧬 Glioma Stage Classifier Model
+# 🧬 Glioma Stage Prediction Model (MATCHES `glioma_stages.pth`)
 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)
-        )
+        self.fc1 = nn.Linear(9, 128)
+        self.relu1 = nn.ReLU()
+        self.fc2 = nn.Linear(128, 64)
+        self.relu2 = nn.ReLU()
+        self.fc3 = nn.Linear(64, 32)
+        self.relu3 = nn.ReLU()
+        self.out = nn.Linear(32, 4)
 
     def forward(self, x):
-        return self.model(x)
+        x = self.relu1(self.fc1(x))
+        x = self.relu2(self.fc2(x))
+        x = self.relu3(self.fc3(x))
+        return self.out(x)