diff --git a/train.py b/train.py
index 878c9a07f560..460c6cd77a0d 100644
--- a/train.py
+++ b/train.py
@@ -419,7 +419,8 @@ def train(hyp,  # path/to/hyp.yaml or hyp dictionary
                                             save_json=is_coco,
                                             verbose=True,
                                             plots=True,
-                                            callbacks=callbacks)  # val best model with plots
+                                            callbacks=callbacks,
+                                            compute_loss=compute_loss)  # val best model with plots
 
         callbacks.run('on_train_end', last, best, plots, epoch)
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")