Can't import torch on my machine with CUDA 12.4

licecolony: <p>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.</p>
<p>I’m on Ubuntu 22.04 with <code>python 3.10.12</code>.<br>
I installed <code>torch</code> via <code>pip3 install torch torchvision torchaudio</code></p>
<p>If I run python3:</p>
<pre><code class="lang-auto">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.
&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.version.cuda
'12.1'
</code></pre>
<p>When I run <code>nvcc -V</code> my output is:</p>
<pre><code class="lang-auto">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
</code></pre>
<p>When I run <code>nvidia-smi</code> my output is:</p>
<pre><code class="lang-auto">NVIDIA-SMI 550.67                 Driver Version: 550.67         CUDA Version: 12.4 
</code></pre>
<p>I have an NVIDIA GeForce RTX 3050 Ti.<br>
Based on Table 3 of <a href="https://docs.nvidia.com/deploy/cuda-compatibility/index.html" class="inline-onebox" rel="noopener nofollow ugc">CUDA Compatibility :: NVIDIA GPU Management and Deployment Documentation</a>, CUDA 12.4 seems like the right version for my NVIDIA driver.</p>
<p>I’m able to run <code>python3 -c 'import torch'</code> with no output, which I assume is good news.<br>
That being said, when I try to import torch into a jupyter notebook, I get the error:</p>
<pre><code class="lang-auto">ModuleNotFoundError: No module named 'torch._custom_ops'; 'torch' is not a package
</code></pre>
<p>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?</p>
<p>I found this: <a href="https://stackoverflow.com/questions/77498488/loading-a-pretrained-model-from-torch-hub-in-sagemaker" class="inline-onebox" rel="noopener nofollow ugc">python - Loading a pretrained model from torch.hub in Sagemaker - Stack Overflow</a> 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.</p>
<p>I would appreciate any help, insight, or simply comments telling me a better place to be asking this question.<br>
Thank you</p>

ptrblck: <aside class="quote no-group" data-username="licecolony" data-post="1" data-topic="202224">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/licecolony/48/69811_2.png" class="avatar"> licecolony:</div>
<blockquote>
<p>I’m able to run <code>python3 -c 'import torch'</code> with no output, which I assume is good news.<br>
That being said, when I try to import torch into a jupyter notebook, I get the error:</p>
</blockquote>
</aside>
<p>It seems your Jupter environment does not use the same PyTorch binary as your terminal, so you would need to fix it by e.g. reinstalling PyTorch into the same env of your Jupyter.</p>
<p>I don’t see any issues pointing towards CUDA as a native PyTorch op import is already failing.</p>