Python PyTorch sharded_model_copy用法及代码示例

本文简要介绍python语言中 torchrec.distributed.utils.sharded_model_copy 的用法。

用法: class torchrec.distributed.utils.sharded_model_copy(device: Optional[Union[str, int, torch.device]])

基础：object

允许将 DistributedModelParallel 模块复制到目标设备。

例子：

# Copying model to CPU.

m = DistributedModelParallel(m)
with sharded_model_copy("cpu"):
        m_cpu = copy.deepcopy(m)

相关用法

注：本文由纯净天空筛选整理自pytorch.org大神的英文原创作品 torchrec.distributed.utils.sharded_model_copy。非经特殊声明，原始代码版权归原作者所有，本译文未经允许或授权，请勿转载或复制。