Update app_flash.py

app_flash.py

@@ -23,14 +23,17 @@ print(f"🔧 Using device: {device} (CPU-only)")
 # 1️⃣ FlashPack model with better hidden layers
 # ============================================================
 class GemmaTrainer(nn.Module, FlashPackMixin):
-    def __init__(self, input_dim: int, hidden_dim: int = 1024, output_dim: int = 1536):
+    def __init__(self):
         super().__init__()
+        input_dim = 1536
+        hidden_dim = 1024
+        output_dim = 1536
         self.fc1 = nn.Linear(input_dim, hidden_dim)
         self.relu = nn.ReLU()
         self.fc2 = nn.Linear(hidden_dim, hidden_dim)
         self.fc3 = nn.Linear(hidden_dim, output_dim)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor):
         x = self.fc1(x)
         x = self.relu(x)
         x = self.fc2(x)