@Adam_Dziedzic I think there is an issue open for that feature    I am pretty sure we would be happy to accept a PR adding this !
Edit: looking at it, it is actually issue number 88, quite an old one  https://github.com/pytorch/pytorch/issues/88 2