diff --git a/pytorch_transformers/modeling_utils.py b/pytorch_transformers/modeling_utils.py
index 5a89badb..c69cba49 100644
--- a/pytorch_transformers/modeling_utils.py
+++ b/pytorch_transformers/modeling_utils.py
@@ -379,11 +379,15 @@ class PreTrainedModel(nn.Module):
                 for head in heads:
                     if head not in self.config.pruned_heads[int(layer)]:
                         self.config.pruned_heads[int(layer)].append(head)
-                        to_be_pruned[int(layer)].append(head)
+
+                        if int(layer) in to_be_pruned:
+                            to_be_pruned[int(layer)].append(head)
+                        else:
+                            to_be_pruned[int(layer)] = [head]
                     else:
-                        logger.warning("Tried to remove head " + head +
-                                       " of layer " + layer +
-                                       " but it was already removed. The current removed heads are " + heads_to_prune)
+                        logger.warning("Tried to remove head " + str(head) +
+                                       " of layer " + str(layer) +
+                                       " but it was already removed. The current removed heads are " + str(heads_to_prune))
 
         base_model._prune_heads(to_be_pruned)