Update watchdog.py

watchdog.py

@@ -85,10 +85,16 @@ def retrain_model():
         if input_ids is None:
             return "⚠️ Not enough data to retrain.", None, "Please log more feedback first."
 
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
         config = mutate_config()
-        model = EvoTransformerForClassification(config)
+        model = EvoTransformerForClassification(config).to(device)
         model.train()
 
+        input_ids = input_ids.to(device)
+        attention_masks = attention_masks.to(device)
+        labels = labels.to(device)
+
         optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
         loss_fn = torch.nn.CrossEntropyLoss()
 
@@ -100,6 +106,10 @@ def retrain_model():
             optimizer.step()
             print(f"🔁 Epoch {epoch+1}: Loss = {loss.item():.4f}")
 
+        # Sanity check logits
+        if logits.shape[-1] < 2:
+            raise ValueError("Logits shape invalid. Retrained model did not output 2 classes.")
+
         # Accuracy
         model.eval()
         with torch.no_grad():
@@ -126,8 +136,9 @@ def retrain_model():
         with open(log_path, "w") as f:
             json.dump(history, f, indent=2)
 
-        # Save model
+        # Save model + tokenizer
         model.save_pretrained("trained_model")
+        tokenizer.save_pretrained("trained_model")
         print("✅ EvoTransformer retrained and saved.")
 
         # Load updated summary + plot