ModuleNotFoundError: No module named 'torch._prims_common'

username: amulya_rao
cooked: <p>Hello. I am new to using Pytorch. I am setting up yolo nas for deepstream as per marcoslucianops deepstream yolo repo for yolo nas.<br>
While generating the onnx model (python3 export_yolonas.py -m yolo_nas_s -w yolo_nas_s_coco.pth --dynamic), I get the following error.</p>
<p>$python3 export_yolonas.py -m yolo_nas_s -w yolo_nas_s_coco.pth --dynamic<br>
Traceback (most recent call last):<br>
File “export_yolonas.py”, line 7, in <br>
import torch.nn as nn<br>
File “/opt/nvidia/deepstream/deepstream-6.2/.venv/lib/python3.8/site-packages/torch/nn/<strong>init</strong>.py”, line 1, in <br>
from .modules import *  # noqa: F403<br>
File “/opt/nvidia/deepstream/deepstream-6.2/.venv/lib/python3.8/site-packages/torch/nn/modules/<strong>init</strong>.py”, line 1, in <br>
from .module import Module<br>
File “/opt/nvidia/deepstream/deepstream-6.2/.venv/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 8, in <br>
from torch._prims_common import DeviceLikeType<br>
ModuleNotFoundError: No module named ‘torch._prims_common’</p>
<p>The torch details are as below<br>
Name: torch<br>
Version: 2.3.1<br>
Summary: Tensors and Dynamic neural networks in Python with strong GPU acceleration<br>
Home-page: <a href="https://pytorch.org/" rel="noopener nofollow ugc">https://pytorch.org/</a><br>
Author: PyTorch Team<br>
Author-email: <a href="mailto:packages@pytorch.org">packages@pytorch.org</a><br>
License: BSD-3<br>
Location: /opt/nvidia/deepstream/deepstream-6.2/.venv/lib/python3.8/site-packages<br>
Requires: filelock, fsspec, jinja2, networkx, nvidia-cublas-cu12, nvidia-cuda-cupti-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-runtime-cu12, nvidia-cudnn-cu12, nvidia-cufft-cu12, nvidia-curand-cu12, nvidia-cusolver-cu12, nvidia-cusparse-cu12, nvidia-nccl-cu12, nvidia-nvtx-cu12, sympy, triton, typing-extensions<br>
Required-by: data-gradients, imagededup, super-gradients, torchaudio, torchmetrics, torchvision</p>

username: ptrblck
cooked: <p>I cannot reproduce the issue using <code>torch==2.3.1</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">python -c "import torch; print(torch.__version__); from torch._prims_common import DeviceLikeType; print(DeviceLikeType)"
2.3.1+cu121
typing.Union[str, torch.device, int]
</code></pre>

username: sumair_arri
cooked: <p>still getting the same error</p>

username: GU_ZHONGXIANG
cooked: <p>Hello Amulya!</p>
<p>Here I propose a possible cause of your problem and its solution:</p>
<p>This error could have been caused by an incomplete install or uninstall of your torch package. To check, see if there is any directory called “~orch” in your python path (/opt/nvidia/deepstream/deepstream-6.2/.venv/lib/python3.8/site-packages). If there it is, simply delete the directory and uninstall and reinstall your torch package and the missing modules should be there in the new installation.</p>
<p>This worked for me.</p>

username: amulya_rao
cooked: <p>I did this and it worked! Thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: lilamonroe
cooked: <p>Thanks for this solution, this is very informative.</p>

username: Ashish_Omar
cooked: <p>Yes<br>
pip3 uninstall torch<br>
pip3 install torch</p>
<p>It worked!</p>