Bugfix in RoBERTa (#241)

* Bugfix in RoBERTa * Add NameError check
asyml · Nov 2, 2019 · 4fee264 · 4fee264
1 parent 9cafd4d
commit 4fee264
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/texar/torch/modules/pretrained/roberta.py b/texar/torch/modules/pretrained/roberta.py
@@ -198,7 +198,7 @@ def _init_from_checkpoint(self, pretrained_model_name: str,
                 pointer.data = tensor.data.type(pointer.dtype)
             elif name.startswith('decoder.sentence_encoder.layers.'):
                 name = name.lstrip('decoder.sentence_encoder.layers.')
-                layer_num, layer_name = name[0], name[2:]
+                layer_num, layer_name = name.split('.', 1)
                 if layer_name in attention_tensor_map:
                     v_names = attention_tensor_map[layer_name]
                     if isinstance(v_names, str):
@@ -214,3 +214,5 @@ def _init_from_checkpoint(self, pretrained_model_name: str,
                                 v_names[i].format(layer_num))
                             assert pointer.shape == tensors[i].shape
                             pointer.data = tensors[i].data.type(pointer.dtype)
+                else:
+                    raise NameError(f"Layer name '{layer_name}' not found")