First, i want to know my module’s location in libtorch like cpu or cuda:0.
Is there any api in libtorch?
torch::jit::module model = torch::jit::load(model_path);
model.??? <- ??? is the api that prints the model location (device)
Second, Can i directly copy the module from gpu:0 to gpu:1, using NVLinks?
torch::Device device = torch::Device(torch::kCUDA, 0);
torch::jit::module model_a = torch::jit::load(model_path);
model_a.to(device)
torch::jit::module model_b = ??? <- ??? is the api that copies the model directly using NVLinks. do not across the host dram.
A model itself does not belong to any specific device as its parameters and buffers can be located on different devices. Assuming you’ve moved the parameters and buffers to a single device, you could grab e.g. the first parameter and check its device attribute instead.
NVLink will be automatically used when detected and you won’t be able to enable or disable it from your SW stack.