Merge pull request #211 from HazyResearch/fix_200

Update where output_dim is defined
HazyResearch · Sep 4, 2019 · b54631e · b54631e
2 parents 8093115 + 0d995b0
commit b54631e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/metal/multitask/mt_end_model.py b/metal/multitask/mt_end_model.py
@@ -98,7 +98,6 @@ def _build(self, input_modules, middle_modules, head_modules):
 
     def _build_input_layer(self, input_modules):
         if input_modules is None:
-            output_dim = self.config["layer_out_dims"][0]
             input_modules = IdentityModule()
 
         if isinstance(input_modules, list):
@@ -107,6 +106,7 @@ def _build_input_layer(self, input_modules):
                 for mod in input_modules
             ]
         else:
+            output_dim = self.config["layer_out_dims"][0]
             input_layer = self._make_layer(
                 input_modules,
                 "input",