From 531cff509dbd82ec5def5279f212689d9547a147 Mon Sep 17 00:00:00 2001
From: tian <11429339@qq.com>
Date: Wed, 10 Jun 2026 09:41:19 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E7=BC=BAPyTorch=E6=97=B6=E8=AE=AD?=
 =?UTF-8?q?=E7=BB=83=E4=B8=8D=E4=BC=9A=E5=9B=A0ImportError=E4=B8=AD?=
 =?UTF-8?q?=E6=96=AD=E6=95=B4=E4=B8=AA=E6=B5=81=E7=A8=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/model_trainer.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/model_trainer.py b/src/model_trainer.py
index 71385dd..12769ed 100644
--- a/src/model_trainer.py
+++ b/src/model_trainer.py
@@ -500,11 +500,15 @@ class ModelTrainer:
                     model.fit(X_train, y_train)
                     
                 elif model_type == 'pytorch':
-                    # 训练PyTorch模型
-                    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:
+                    # 训练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
                 
                 # 评估模型性能