本文整理汇总了Python中torch.cuda.comm.scatter方法的典型用法代码示例。如果您正苦于以下问题：Python comm.scatter方法的具体用法？Python comm.scatter怎么用？Python comm.scatter使用的例子？那么恭喜您, 这里精选的方法代码示例或许可以为您提供帮助。您也可以进一步了解该方法所在类torch.cuda.comm的用法示例。
在下文中一共展示了comm.scatter方法的1个代码示例，这些例子默认根据受欢迎程度排序。您可以为喜欢或者感觉有用的代码点赞，您的评价将有助于系统推荐出更棒的Python代码示例。
示例1: forward
# 需要导入模块: from torch.cuda import comm [as 别名]
# 或者: from torch.cuda.comm import scatter [as 别名]
def forward(ctx, target_gpus, chunk_sizes, dim, input):
ctx.target_gpus = target_gpus
ctx.chunk_sizes = chunk_sizes
ctx.dim = dim
ctx.input_device = input.get_device() if input.is_cuda else -1
streams = None
if ctx.input_device == -1:
# Perform CPU to GPU copies in a background stream
streams = [_get_stream(device) for device in ctx.target_gpus]
outputs = comm.scatter(input, ctx.target_gpus, ctx.chunk_sizes, ctx.dim, streams)
# Synchronize with the copy stream
if streams is not None:
for i, output in enumerate(outputs):
with torch.cuda.device(ctx.target_gpus[i]):
main_stream = torch.cuda.current_stream()
main_stream.wait_stream(streams[i])
output.record_stream(main_stream)
return outputs