layer drop
Summary: This diff enables layer drop in transformer decoder in production training pipeline (ptt_transformer). It builds on top of the fairseq implementation added by Angela Fan, and added additional logic to handle corresponding dropping layers at test time in exported model.
@@ -345,7 +345,7 @@ def prune_state_dict(state_dict, args):
It's called by functions that load models from checkpoints and does not
need to be called directly.
"""
- if not args:
+ if not args or args.arch == "ptt_transformer":
# args should not be none, but don't crash if it is.
return state_dict
