Fix BatchNorm1d error: enable only Dropout (not BatchNorm) for MC uncertainty

app.py

@@ -84,6 +84,20 @@ NETWORK_FUNCTIONS = {
 }
 
 
+# ============================================================
+# Helper: enable only Dropout for MC sampling (keep BatchNorm in eval)
+# ============================================================
+def enable_dropout_only(model):
+    """Enable Dropout layers while keeping BatchNorm in eval mode.
+    
+    This is needed for MC Dropout uncertainty estimation with batch_size=1,
+    because BatchNorm1d requires batch_size > 1 in training mode.
+    """
+    for module in model.modules():
+        if isinstance(module, nn.Dropout):
+            module.train()
+
+
 # ============================================================
 # BrainEncoder model (must match training architecture exactly)
 # ============================================================
@@ -412,18 +426,22 @@ class ModelManager:
             with torch.no_grad():
                 deep_pred, intermediates = self.brain_encoder(input_features, return_intermediates=True)
             
-            # Compute uncertainty via dropout MC (deep encoder)
-            self.brain_encoder.train()
+            # Compute uncertainty via MC Dropout
+            # IMPORTANT: Only enable Dropout layers, keep BatchNorm in eval mode.
+            # BatchNorm1d requires batch_size > 1 in training mode, but we have batch_size=1.
+            self.brain_encoder.eval()  # Ensure everything is in eval mode first
+            enable_dropout_only(self.brain_encoder)  # Selectively enable only Dropout
             mc_predictions = []
             for _ in range(10):
                 with torch.no_grad():
                     mc_pred = self.brain_encoder(input_features)
                     mc_predictions.append(mc_pred.cpu().numpy().flatten())
-            self.brain_encoder.eval()
+            self.brain_encoder.eval()  # Restore full eval mode
             mc_predictions = np.array(mc_predictions)
             uncertainty = np.std(mc_predictions, axis=0)
         else:
             # Estimate uncertainty from ridge prediction variance across feature perturbation
+            ridge = self.ridge_model
             mc_predictions = []
             for _ in range(10):
                 noise = np.random.normal(0, 0.01, size=input_features_np.shape)