Similar to https://github.com/Adversarial-Deep-Learning/code-soup/issues/97. Reference: https://github.com/QData/TextAttack/blob/master/textattack/models/wrappers/pytorch_model_wrapper.py
