colossalai.zero.shard_utils.commons

colossalai.zero.shard_utils.commons.get_shard(tensor, rank, world_size)[source]

Return the local shard of a full tensor.