C++ Multiple GPUs for inference

Hi Everyone,

I am unable to find any documentation on how to set multiple GPUs for inference.

In python the following can be done:

device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)

model = CreateModel()

model= nn.DataParallel(model)

However for C++ I can’t find the equivalent or any documentation.
torch::nn does not contain either DataParallel or DistributedDataParallel

Any help is appreciated, thanks!

libtorch uses the parallel::data_parallel API as seen e.g. here. I don’t know if there are other wrappers, but based on the examples I would guess they should work in a similar way as the DataParallel wrapper in Python.

Thank you for the snippet!!

I forgot to include

#include <torch/nn.h>
#include <torch/nn/parallel/data_parallel.h>