torchfoo.distributed.all_concat_jagged

torchfoo.distributed.all_concat_jagged(x)

Concatenate a Tensor with varying dim 0 from multiple GPUs

Not differentiable.

Parameters:

x (Tensor) – a Tensor on this device

Returns:

Concatenated tensor from all processes along dim 0.

Return type:

Tensor