fix: PyTorch训练时未传equipment_type导致火箭炮模型架构错误

2026-06-10 11:00:55 +08:00 · 2026-06-10 11:00:55 +08:00 · 42891d2980
commit 42891d2980
parent 2809631e87
1 changed files with 1 additions and 1 deletions
--- a/src/model_trainer.py
+++ b/src/model_trainer.py
@ -522,7 +522,7 @@ class ModelTrainer:
                    try:
                        train_dataset = EquipmentDataset(X_train, y_train)
                        train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
-                        training_success = self.train_model(train_loader)
+                        training_success = self.train_model(train_loader, equipment_type=equipment_type)
                        if not training_success:
                            continue
                    except ImportError as e: