ModuleNotFoundError: No module named 'torch._custom_ops'; 'torch' is not a package

Hi there, I’m not sure if this is the right place to ask, but I just installed cuda tools to run some GPU-based machine learning stuff on my computer, and I’m running into an issue importing torch.

I’m on Ubuntu 22.04 with python 3.10.12

If I run python3:

Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.version.cuda
'12.1'

When I run nvcc -V my output is:

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Thu_Mar_28_02:18:24_PDT_2024
Cuda compilation tools, release 12.4, V12.4.131
Build cuda_12.4.r12.4/compiler.34097967_0

When I run nvidia-smi my output is:

NVIDIA-SMI 550.67                 Driver Version: 550.67         CUDA Version: 12.4 

I have an NVIDIA GeForce RTX 3050 Ti.
Based on Table 3 of CUDA Compatibility :: NVIDIA GPU Management and Deployment Documentation, CUDA 12.4 seems like the right version for my NVIDIA driver.

I’m able to run python3 -c 'import torch' with no output, which I assume is good news.
That being said, when I try to import torch into a jupyter notebook, I get the error:

ModuleNotFoundError: No module named 'torch._custom_ops'; 'torch' is not a package

I was able to find torch._custom_ops myself, so I know it exists, but I’m not sure why it isn’t working in Jupyter Notebook?

I found this: python - Loading a pretrained model from torch.hub in Sagemaker - Stack Overflow but that didn’t seem relevant given I’m not using Sagemaker and simply trying to get my local machine ready to tackle GPU training tasks.

I would appreciate any help, insight, or simply comments telling me a better place to be asking this question.
Thank you

2 Likes

I have same problem too.Please help us guys

2 Likes

I follow the post to install Torch but got errors
PyTorch for Jetson

ERROR: pip’s dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
torchvision 0.16.0 requires torch==2.1.0, but you have torch 2.1.0a0+41361538.nv23.6 which is incompatible.

And got errors when using torch module

BRuntimeError: Couldn’t load custom C++ ops. This can happen if your PyTorch and torchvision versions are incompatible, or if you had errors while compiling torchvision from source. For further information on the compatible versions, check GitHub - pytorch/vision: Datasets, Transforms and Models specific to Computer Vision for the compatibility matrix. Please check your PyTorch version with torch.version and your torchvision version with torchvision.version and verify if they are compatible, and if not please reinstall torchvision so that it matches your PyTorch install.

After I uninstalled PyTorch and Torchvision and reinstall PyTorch for Jetson and Trochvidsion 0.16.1 from source code Compile TorchVision 0.16.1, it works now.