Further fix for encoder_only (#1242)

nyu-mll · Dec 3, 2020 · 50b0116 · 50b0116
1 parent 1f66050
commit 50b0116
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jiant/proj/main/modeling/model_setup.py b/jiant/proj/main/modeling/model_setup.py
@@ -215,7 +215,7 @@ def load_encoder_only(jiant_model, weights_dict):
     for taskmodel_key in jiant_model.task_to_taskmodel_map:
         for encoder_key in encoder_keys:
             new_key = f"taskmodels_dict.{taskmodel_key}.encoder.{encoder_key}"
-            new_weights_dict[new_key] = weights_dict[new_key]
+            new_weights_dict[new_key] = weights_dict[f"encoder.{encoder_key}"]
 
     mismatch = jiant_model.load_state_dict(new_weights_dict, strict=False)
     assert not mismatch.unexpected_keys