Pytorch C ++-在Stream中拥有内核计算

Question

我在Pytorch Python中有以下几点：

torch.cuda.synchronize()
with torch.cuda.stream(s1):
    A = torch.mm(X, Y)
with torch.cuda.stream(s2):
    B = bigCpuOperation().cuda()
# Wait for A and B to be computed.
torch.cuda.synchronize()
# Do stuff with A and B

我需要用Pytorch C ++实现相同的功能。我已经实例化了流，但是我没有看到执行上述操作的机制。例如torch::mm不允许我通过流