@aclifton314 You can perform generic calculations in pytorch using multiple gpus similar to the code example you provided. Basically spawn multiple processes where each process drives a single GPU and have each GPU do part of the computation. Then you can use PyTorch collective APIs 4 to perform any aggregations across GPUs that you need.