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。非經特殊聲明，原始代碼版權歸原作者所有，本譯文未經允許或授權，請勿轉載或複製。