Pytorch train on cpu works but not on gpu

Kennyruffles: <p>Hi,</p>
<p>I’ve installed torch using conda on Windows 10 (conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch)</p>
<p>python</p>
<blockquote>
<blockquote>
<blockquote>
<p>import torch<br>
torch.cuda.is_available()<br>
True<br>
torch.cuda.get_device_name(0)<br>
‘NVIDIA GeForce GTX 1650’</p>
</blockquote>
</blockquote>
</blockquote>
<blockquote>
<blockquote>
<blockquote>
<p>import torch<br>
print(torch.<strong>version</strong>)<br>
1.10.0</p>
</blockquote>
</blockquote>
</blockquote>
<blockquote>
<blockquote>
<blockquote>
<p>import numpy<br>
numpy.version.version<br>
‘1.21.2’</p>
</blockquote>
</blockquote>
</blockquote>
<p>When running train.py file for yolov5 as:<br>
python train.py --img 512 --workers 1 --batch 2 --epochs 20 --data yolo_train.yaml --weights yolov5s.pt --cache</p>
<p>It shows the following message:<br>
YOLOv5  2021-11-4 torch 1.10.0 CUDA:0 (NVIDIA GeForce GTX 1650, 4096MiB)<br>
0/19    0.325G       nan       nan       nan         2       512: 100%|████████████████████████████████| 152/152 [00:40&lt;00:00,  3.78it/s]<br>
C:\Users\m.conda\envs\yolov5_gpu3\lib\site-packages\torch\optim\lr_scheduler.py:129: UserWarning: Detected call of <code>lr_scheduler.step()</code> before <code>optimizer.step()</code>. In PyTorch 1.1.0 and later, you should call them in the opposite order: <code>optimizer.step()</code> before <code>lr_scheduler.step()</code>.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at <a href="https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate" class="inline-onebox" rel="noopener nofollow ugc">torch.optim — PyTorch 2.1 documentation</a><br>
warnings.warn("Detected call of <code>lr_scheduler.step()</code> before <code>optimizer.step()</code>. "<br>
Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|████████████████| 424/424 [00:52&lt;00:00,  8.06it/s]<br>
all       1694          0          0          0          0          0</p>
<pre><code> Epoch   gpu_mem       box       obj       cls    labels  img_size
  1/19    0.357G       nan       nan       nan         4       512: 100%|████████████████████████████████| 152/152 [00:35&lt;00:00,  4.26it/s]
           Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|████████████████| 424/424 [00:52&lt;00:00,  8.04it/s]
             all       1694          0          0          0          0          0

 Epoch   gpu_mem       box       obj       cls    labels  img_size
  2/19    0.357G       nan       nan       nan         2       512: 100%|████████████████████████████████| 152/152 [00:35&lt;00:00,  4.28it/s]
</code></pre>
<p>Basically,  it never learns. Box, obj, cls are always nan, and P and R are always 0.</p>
<p>On the other hand, when using the same code but running on cpu:</p>
<p>Logging results to runs\train\exp25<br>
Starting training for 20 epochs…</p>
<pre><code> Epoch   gpu_mem       box       obj       cls    labels  img_size
  0/19        0G    0.1147   0.02892   0.05367         5       512:   2%|▋                                 | 3/152 [00:05&lt;03:40,  1.48s/it]
</code></pre>
<p>The values are correct.</p>
<p>Do you have an idea how to solve that?</p>
<p>Thanks</p>