diff --git a/python-lib/dku_deeplearning_image/tensorboard_handler.py b/python-lib/dku_deeplearning_image/tensorboard_handler.py
index 2bd4ef8..af1a5ad 100644
--- a/python-lib/dku_deeplearning_image/tensorboard_handler.py
+++ b/python-lib/dku_deeplearning_image/tensorboard_handler.py
@@ -37,6 +37,7 @@ def download_logs_to_local(folder):
 
 
 def __get_logs_path():
+
     retrained_model_folder = dataiku.Folder(get_webapp_config().get('retrained_model_folder'))
     check_folder_exists(retrained_model_folder)
 
diff --git a/webapps/tensorboard-web-app-v2/webapp.json b/webapps/tensorboard-web-app-v2/webapp.json
index 9438fa1..b0ad8eb 100644
--- a/webapps/tensorboard-web-app-v2/webapp.json
+++ b/webapps/tensorboard-web-app-v2/webapp.json
@@ -6,16 +6,17 @@
     },
     "libraries" : [],
     "hasBackend": "true",
-    "paramsPythonSetup": "helpers/webapp-helper.py",
     "params": [
         {
             "name": "retrained_model_folder",
-            "type": "SELECT",
+            "type": "MANAGED_FOLDER",
             "label": "Folder containing retrained model",
-            "mandatory": true,
-            "getChoicesFromPython": true
+            "mandatory": true
         }
     ],
     "baseType": "STANDARD",
-    "roles": []
+    "roles": [{
+        "type": "MANAGED_FOLDER",
+        "targetParamsKey": "retrained_model_folder"
+    }]
 }