Python PyTorch stash用法及代碼示例

本文簡要介紹python語言中 torch.distributed.pipeline.sync.skip.skippable.stash 的用法。

用法:

class torch.distributed.pipeline.sync.skip.skippable.stash(name, tensor)

參數：

• name(str) -跳躍張量的名稱

• input(torch.Tensor或者None) -傳遞給跳躍連接的張量

存儲跳過張量的命令。

def forward(self, input):
    yield stash('name', input)
    return f(input)