Process killed when loading cpp extensions

1

I was trying to load a cpp extension with torch.utils.cpp_extension.load() on my Jetson TX2 device with CUDA-10.2.

from torch.utils.cpp_extension import load
matrix_op_cuda = load(
    'matrix_op_cuda', ['./matrix_op/matrix_op_cuda.cpp', './matrix_op/matrix_op_cuda_kernel.cu'], verbose=True)

The compilation process with ninja went smoothly, but when loading the extension module, the process was killed and the log showed no more info.

I tried to load the extension again but the same thing happened.

Building extension module matrix_op_cuda...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module matrix_op_cuda...
Killed

Wonder what could be the problem?

2

Killed is usually raised if the OS sends a SIGKILL to the process if you are running out of memory. Based on the output it seems that the compilation succeeds but the loading of the kernel is raising the error? Try to check the memory usage during the call and see if you are indeed running OOM.