fix: 缺PyTorch时训练不会因ImportError中断整个流程

2026-06-10 09:41:19 +08:00 · 2026-06-10 09:41:19 +08:00 · 531cff509d
commit 531cff509d
parent b6b072d844
1 changed files with 9 additions and 5 deletions
--- a/src/model_trainer.py
+++ b/src/model_trainer.py
@ -500,12 +500,16 @@ class ModelTrainer:
                    model.fit(X_train, y_train)
                    
                elif model_type == 'pytorch':
-                    # 训练PyTorch模型
+                    # 训练PyTorch模型（如未安装则跳过）
+                    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)
                        if not training_success:
                            continue
+                    except ImportError as e:
+                        logger.warning(f"PyTorch not available, skipping: {e}")
+                        continue
                
                # 评估模型性能
                if model_type == 'pytorch':