diff --git a/src/model_trainer.py b/src/model_trainer.py
index 70b9cf9..7d7dfc5 100644
--- a/src/model_trainer.py
+++ b/src/model_trainer.py
@@ -252,13 +252,16 @@ class ModelTrainer:
             
             # 保存模型
             model_path = f'{model_dir}/{equipment_type_en}_{timestamp}.pth'
-            torch.save({
+            checkpoint = {
                 'model_state_dict': self.model.state_dict(),
-                'input_size': self.model.equipment_net[0].in_features + 5,
-                'manufacturer_net_state': self.model.manufacturer_net.state_dict(),
-                'equipment_net_state': self.model.equipment_net.state_dict(),
-                'combined_net_state': self.model.combined_net.state_dict()
-            }, model_path)
+                'equipment_type': equipment_type
+            }
+            # input_size：火箭炮取 net[0].in_features，巡飞弹取 equipment_net[0].in_features+5
+            if equipment_type == '火箭炮':
+                checkpoint['input_size'] = self.model.net[0].in_features
+            else:
+                checkpoint['input_size'] = self.model.equipment_net[0].in_features + 5
+            torch.save(checkpoint, model_path)
             
             # 保存标准化器
             scaler_path = f'{model_dir}/{equipment_type_en}_{timestamp}_scaler.pth'
@@ -329,13 +332,24 @@ class ModelTrainer:
                 
                 # 加载模型
                 if model_record['model_type'] == 'pytorch':
-                    # 加载PyTorch模型
                     checkpoint = torch.load(model_record['model_path'], encoding='latin1')
-                    input_size = checkpoint['input_size']
+                    input_size = checkpoint.get('input_size')
                     
                     # 创建新模型实例
                     self.model = CostPredictionModel(input_size, equipment_type).to(self.device)
-                    self.model.load_state_dict(checkpoint['model_state_dict'])
+                    
+                    # 尝试加载 state_dict，架构不匹配则标记模型为无效
+                    try:
+                        self.model.load_state_dict(checkpoint['model_state_dict'])
+                    except RuntimeError as e:
+                        logger.error(f"Architecture mismatch for {model_record['model_name']}: {e}")
+                        # 标记此模型为非激活，下次不再尝试
+                        conn.cursor().execute(
+                            "UPDATE trained_models SET is_active = 0 WHERE id = ?",
+                            (model_record['id'],)
+                        )
+                        conn.commit()
+                        raise ValueError(f"模型架构已更新，请重新训练 {equipment_type} 的 PyTorch 模型")
                     
                     # 加载标准化器
                     scalers = torch.load(model_record['scaler_path'], encoding='latin1')