pytorch

d316a2dd - [spmd] Enable data parallel to work with non 0 batch dim (#100073)

[spmd] Enable data parallel to work with non 0 batch dim (#100073) This PR enables data parallel to work with non 0 batch dim, the only thing we need to do is to expose the input_batch_dim to DataParallelMode and the data parallel expansion automatically works as we have done things correctly in batch dim analysis. Pull Request resolved: https://github.com/pytorch/pytorch/pull/100073 Approved by: https://github.com/mrshenli