diff --git a/src/mattersim/utils/mlflow_utils.py b/src/mattersim/utils/mlflow_utils.py
index 3d07759..d97b1bf 100644
--- a/src/mattersim/utils/mlflow_utils.py
+++ b/src/mattersim/utils/mlflow_utils.py
@@ -82,8 +82,9 @@ def predict(self, context, model_input, params=None):
         with tempfile.TemporaryDirectory() as td:
             data["work_dir"] = td
 
-        df = wrapper(atoms_list, **data)
-        output = json.loads(df.to_json(orient="records"))
+            df = wrapper(atoms_list, **data)
+            output = json.loads(df.to_json(orient="records"))
+
         # print(f"Output:\n{output}")
         return output