one more thing

eflows4hpc · Feb 5, 2024 · 0774b1a · 0774b1a
1 parent d9110c7
commit 0774b1a
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/model_search_upload.py b/model_search_upload.py
@@ -94,13 +94,12 @@ def uploat_to_mlflow(temp_dir, **context):
                 for m in metrics:
                     if m not in dct:
                         continue
-                    
+
                     print(f"Logging metric {m} {dct[m][i]}")
                     mlflow.log_metric(key=m, value=dct[m][i])
 
         #clean up
-        if 'temp_dir' in attrs:
-            shutil.rmtree(path=attrs["temp_dir"])
+        shutil.rmtree(temp_dir)
 
     setup_task = PythonOperator(python_callable=setup, task_id="setup_connection")
     a_id = setup_task.output["return_value"]