Hi @ch283 ,
because you are working with Pytorch, maybe this Feature Request could be of interest.
Regards, Andrés.