torch 计算流stream

import torch

# 创建两个流
stream1 = torch.cuda.Stream()
stream2 = torch.cuda.Stream()

# 使用流进行操作
with torch.cuda.stream(stream1):
    tensor1 = torch.ones(10000, device='cuda')
    tensor2 = tensor1 * 2

with torch.cuda.stream(stream2):
    tensor3 = torch.ones(10000, device='cuda')
    tensor4 = tensor3 * 3

# 等待流中的操作完成
stream1.synchronize()
stream2.synchronize()