Yeah true, that's also a bit annoying in PyTorch compared to TF. I remember reading on the PyTorch forums that this was a bit hard to implement for some reason,

Yeah true, that's also a bit annoying in PyTorch compared to TF. I remember reading on the PyTorch forums that this was a bit hard to implement for some reason, but I can't recall any details if there were any.