[two_in_one]

I wonder if this sort of memory management can be made for Pytorch transformers as under the hood optimization.