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)