Force libtorch to use CUDA context

[nachovall]

I’m trying to integrate C++ libtoch to load a model into my application. My application does many CUDA staff before I load the model with libtorch. Thus, the CUDA context has already been created.

For some reason, even the CUDA context has already been created and the calling thread has already a valid context, when I call

torch::jit::script::Module module = torch::jit::load(“test.pt”);
module.to(at::kCUDA);

a new context is created by libtorch. The new context is not event push to the stack of contexts, is overwritting the current context. I know because if after the module.to(at::kCUDA) I call cuCtxPopCurrent the current context is null.

This causes a lot of problems because I cannot interact with current allocated memory I have in my context.

#jit #cuda #c++

cc @suo

[raimis]

We are hitting the same context problem (openmm/openmm-torch#13), while trying to integrate PyTorch with OpenMM (https://github.com/openmm/openmm).

Would it be difficult to modify torch::jit::load to reuse an existing context? Or has somebody found a workaround?