No module named ‘torch._custom_ops’ in Jupyter Notebooks

Question

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 I've tried installing torch within a conda environment and locally. Unfortunately, when I try to import torch into a jupyter notebook, I get the error (Doing literally nothing else in the notebook but importing torch):

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

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 https://docs.nvidia.com/deploy/cuda-compatibility/index.html, 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: Loading a pretrained model from torch.hub in Sagemaker 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 see other posts such as No module named "Torch" pip3 install https://download.pytorch.org/whl/cpu/torch-1.0.1-cp36-cp36m-win_amd64.whl, which I haven't done, but looks like it's for windows. Is that something I need?

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

Answer 1

This is just a version issue...I did

!pip install torch==1.4 torchvision==0.14.1

and it seems to work fine!