diff --git a/python/ray/tune/integration/mlflow.py b/python/ray/tune/integration/mlflow.py
index cbd3811d4..6e038b810 100644
--- a/python/ray/tune/integration/mlflow.py
+++ b/python/ray/tune/integration/mlflow.py
@@ -274,8 +274,8 @@ def mlflow_mixin(func: Callable):
         @mlflow_mixin
         def train_fn(config):
             for i in range(10):
-                loss = self.config["a"] + self.config["b"]
-                mlflow.log_metric(key="loss", value=loss})
+                loss = config["a"] + config["b"]
+                mlflow.log_metric(key="loss", value=loss)
             tune.report(loss=loss, done=True)
 
         tune.run(