diff --git a/src/transformers/integrations.py b/src/transformers/integrations.py
index e11eeb78adabbe..db9a81c0fa38e5 100644
--- a/src/transformers/integrations.py
+++ b/src/transformers/integrations.py
@@ -432,7 +432,7 @@ class MLflowCallback(TrainerCallback):
     MAX_LOG_SIZE = 100
 
     def __init__(self):
-        assert _has_mlflow, "MLflow requires mlflow to be installed. Run `pip install mlflow`."
+        assert _has_mlflow, "MLflowCallback requires mlflow to be installed. Run `pip install mlflow`."
         self._initialized = False
         self._log_artifacts = False