diff --git a/keras-yolo3-master/train.py b/keras-yolo3-master/train.py
index adc2242..cc14742 100644
--- a/keras-yolo3-master/train.py
+++ b/keras-yolo3-master/train.py
@@ -272,11 +272,7 @@ def _main_(args):
         callbacks        = callbacks
     )
 
-
-    # make a GPU version of infer_model for evaluation
-    if multi_gpu > 1:
-        infer_model = load_model(config['train']['saved_weights_name'])
-    infer_model.load_weights(config['train']['saved_weights_name'])
+    train_model.load_weights(config['train']['saved_weights_name'])
     infer_model.save(config['train']['saved_weights_name'])
     ###############################
     #   Run the evaluation