Quick fix for Structured Dropout checkpointing (#1406)

Summary:
Here's a quick fix for https://github.com/pytorch/fairseq/issues/1403.

To keep it short, this fix allows the user to checkpoint a translation model properly after applying layer pruning on a restored transformer file.
Pull Request resolved: https://github.com/pytorch/fairseq/pull/1406

Differential Revision: D18637540

Pulled By: myleott

fbshipit-source-id: 0f5e91e05e6579f0f459bc5293e9b14cb267322d

fairseq/checkpoint_utils.py

@@ -408,6 +408,13 @@ def prune_state_dict(state_dict, args):
                 new_state_key = layer_name[:substitution_match.start(1)] + new_layer_number + layer_name[substitution_match.end(1):]
                 new_state_dict[new_state_key] = state_dict[layer_name]
 
+    # Since layers are now pruned, *_layers_to_keep are no longer needed.
+    # This is more of "It would make it work fix" rather than a proper fix.
+    if "encoder_layers_to_keep" in vars(args):
+        args.encoder_layers_to_keep = None
+    if "decoder_layers_to_keep" in vars(args):
+        args.decoder_layers_to_keep = None
+
     return new_state_dict