How to load a pruned Albert model with from_pretrained()?
# ❓ Questions & Help

<!-- The GitHub issue tracker is primarly intended for bugs, feature requests,
     new models and benchmarks, and migration questions. For all other questions,
     we direct you to Stack Overflow (SO) where a whole community of PyTorch and
     Tensorflow enthusiast can help you out. Make sure to tag your question with the
     right deep learning framework as well as the huggingface-transformers tag: 
     https://stackoverflow.com/questions/tagged/huggingface-transformers 
     
     If your question wasn't answered after a period of time on Stack Overflow, you
     can always open a question on GitHub. You should then link to the SO question 
     that you posted.
     -->

## Details
<!-- Description of your issue -->
I pruned Albert during the fine-tuning phase. I was unable to load the pruned model after saving it. I tried using: 
```
output_model_file = os.path.join(args.output_dir, "pytorch_model.bin")
model_state_dict = torch.load(output_model_file)
model = model_class.from_pretrained(args.output_dir,state_dict=model_state_dict)
```
but still got the same error:
```
 File "run_glue.py", line 526, in main
    model = model_class.from_pretrained(args.output_dir,state_dict=model_state_dict)
  File "/home/user/.local/lib/python3.7/site-packages/transformers/modeling_utils.py", line 471, in from_pretrained
    model.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for AlbertForSequenceClassification:
	size mismatch for albert.encoder.albert_layer_groups.0.albert_layers.0.attention.query.weight: copying a param with shape torch.Size([64, 768]) from checkpoint, the shape in current model is torch.Size([768, 768]).
	size mismatch for albert.encoder.albert_layer_groups.0.albert_layers.0.attention.query.bias: copying a param with shape torch.Size([64]) from checkpoint, the shape in current model is torch.Size([768]).
	size mismatch for albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key.weight: copying a param with shape torch.Size([64, 768]) from checkpoint, the shape in current model is torch.Size([768, 768]).
	size mismatch for albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key.bias: copying a param with shape torch.Size([64]) from checkpoint, the shape in current model is torch.Size([768]).
	size mismatch for albert.encoder.albert_layer_groups.0.albert_layers.0.attention.value.weight: copying a param with shape torch.Size([64, 768]) from checkpoint, the shape in current model is torch.Size([768, 768]).
	size mismatch for albert.encoder.albert_layer_groups.0.albert_layers.0.attention.value.bias: copying a param with shape torch.Size([64]) from checkpoint, the shape in current model is torch.Size([768]).
	size mismatch for albert.encoder.albert_layer_groups.0.albert_layers.0.attention.dense.weight: copying a param with shape torch.Size([768, 64]) from checkpoint, the shape in current model is torch.Size([768, 768]).
```
transformer version == 2.2.1
torch-1.4.0
Anybody can help?



Now I'm using head_mask instead of prune_heads. So, I didn't actually prune heads.