Port fix from https://github.com/pytorch/pytorch/pull/193397
Port fix from pytorch/pytorch#193397