Problem with custom extensions and PyTorch installation using Anaconda

username: MFML
cooked: <p>I installed PyTorch using conda in a fresh environment on my Manjaro linux OS as described in <a href="https://pytorch.org/get-started/locally/" rel="noopener nofollow ugc">“Get Started” / “Start Locally”</a>, using the given command:</p>
<pre><code class="lang-auto">conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
</code></pre>
<p>This worked fine, given that <code>torch.cuda.is_available()</code> returns <code>True</code> and for testing, I successfully trained a model on my GPU (RTX 4090).</p>
<p>However, it looks like I can’t use custom extensions, such as the CUDA kernel provided for <a href="https://github.com/microsoft/Swin-Transformer/tree/main/kernels/window_process" rel="noopener nofollow ugc">swin transformers</a>. If I run the <code>setup.py</code>, I get the error:</p>
<pre><code class="lang-auto">OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.
</code></pre>
<p>I also recognize that <code>nvcc</code> is not available in the shell.</p>
<p>I’ve read quite a lot of postings across different forums which seem to be related, but couldn’t find a solution for this problem.</p>
<p>It looks like I have to install CUDA separately, but I <em>also</em> need to have the version match with the version of PyTorch. If I install CUDA using the package manager of Manjaro (and not via conda), the versions don’t match and it doesn’t work.</p>
<p>How can I install both PyTorch and CUDA <em>with conda</em> to make custom extensions work?</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="MFML" data-post="1" data-topic="188446">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/9d8465/48.png" class="avatar"> MFML:</div>
<blockquote>
<p>It looks like I have to install CUDA separately, but I <em>also</em> need to have the version match with the version of PyTorch.</p>
</blockquote>
</aside>
<p>Yes, that’s correct and you should install a local CUDA toolkit, including the <code>nvcc</code> CUDA compiler, to be able to build custom CUDA extensions.</p>
<aside class="quote no-group" data-username="MFML" data-post="1" data-topic="188446">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/9d8465/48.png" class="avatar"> MFML:</div>
<blockquote>
<p>How can I install both PyTorch and CUDA <em>with conda</em> to make custom extensions work?</p>
</blockquote>
</aside>
<p>You could experiment with the <code>cuda</code> package on <code>conda</code> as it would also ship with <code>nvcc</code>, but I have never used it and don’t know how easy it would be to use compared to just installing CUDA locally.</p>

username: MFML
cooked: <p>Thanks for your answer!</p>
<p>It turned out that installing CUDA in conda worked fine. Just intalled the <a href="https://anaconda.org/nvidia/cuda" rel="noopener nofollow ugc">CUDA package</a> corresponding to the required version number (11.8 in my case), and additionally installed the <code>gcc_linux-64</code> and <code>gxx_linux-64</code> packages, which where needed because the locally installed GCC version wasn’t compatible.</p>
<p>I successfully ran the <code>setup.py</code> script from swin transformer and used the compiled functions.</p>