From 42891d298010fb1795433bcf3283548764e49f7c Mon Sep 17 00:00:00 2001
From: tian <11429339@qq.com>
Date: Wed, 10 Jun 2026 11:00:55 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20PyTorch=E8=AE=AD=E7=BB=83=E6=97=B6?=
 =?UTF-8?q?=E6=9C=AA=E4=BC=A0equipment=5Ftype=E5=AF=BC=E8=87=B4=E7=81=AB?=
 =?UTF-8?q?=E7=AE=AD=E7=82=AE=E6=A8=A1=E5=9E=8B=E6=9E=B6=E6=9E=84=E9=94=99?=
 =?UTF-8?q?=E8=AF=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/model_trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/model_trainer.py b/src/model_trainer.py
index 3fef1e8..bc0b833 100644
--- 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: