diff --git a/combo/combo_model.py b/combo/combo_model.py
index 79c9f36097ab6eeed9c27f1e949eea861c3b868e..d2312319a930ab43b7db1c8cc3a510fdb9a1b775 100644
--- a/combo/combo_model.py
+++ b/combo/combo_model.py
@@ -23,7 +23,7 @@ from combo.nn import utils
 from combo.nn.utils import get_text_field_mask
 from combo.predictors import Predictor
 from combo.utils import metrics
-from utils import ConfigurationError
+from combo.utils import ConfigurationError
 
 
 @Registry.register("semantic_multitask")
diff --git a/combo/main.py b/combo/main.py
index 44a6a30868ea54fac9580881b5f1a1667bfce4c4..055ebeb13a5d6e11c3e4540b3e68e5067ff47902 100755
--- a/combo/main.py
+++ b/combo/main.py
@@ -352,7 +352,11 @@ def _get_ext_vars(finetuning: bool = False) -> Dict:
                 }
             }
         },
-        "pretrained_tokens": FLAGS.pretrained_tokens,
+        "vocabulary": {
+            "parameters": {
+                "pretrained_files": FLAGS.pretrained_tokens
+            }
+        },
         "word_batch_size": int(FLAGS.word_batch_size),
     }