Split tensors as part of split_between_processes

[muellerzr]

Solves #1400 by also letting the tensors be split if added. Note that similar to the list version, the tensors padded will have the value of the last item in the tensor being used. Such that on four processes:

t = torch.tensor([[0,1,2,3], [4,5,6,7]])
with PartialState().split_between_processes(t, with_padding=True) as t:
    print(t)
# Process 1:
tensor([[0,1,2,3]])
# Process 2:
tensor([[4,5,6,7]])
# Process 3:
tensor([[4,5,6,7]])
# Process 4:
tensor([[4,5,6,7]])

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

LGTM, can you add a test of this use case as well?