Which PyTorch version are you using? If an older one, could you update to 1.9.1 or the nightly and rerun your script? We were hitting a missing shape check when calling into cublas via matmul ops, which might cause this downstream error.