Hi,
This should work yes.
But this is not used very often and it appears that we broke it in some very recent versions. The issue tracking the fix is here: https://github.com/pytorch/pytorch/issues/45136 10