How do I place torch.fx.Node on different gpu

torch.compile FX (Functional Transformations)
1

image
image1740×597 29.6 KB

I use torch.fx to trace nn.Module, then I can the node of graph, how do i place this node on different gpu