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: