Could you post the output of python -m torch.utils.collect_env as well as the shapes for all tensors, please?
Based on the error message it could be a valid error, which we would have to debug and fix (if it’s still observed in the current nightly release).