报错：CUDA out of memory

报错：

RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 2.00 GiB total capacity; 1.27 GiB already allocated; 3.24 MiB free; 135.27 MiB cached)

我查了一下发现电脑里面有两个显卡的，默认选择了intel的显卡（即是GPU 0），导致出错。但BATCH_SIZE改为1或者with torch.no_grad()是可以运行的。

后来我改为nvdia显卡，即是GPU 1，BATCH_SIZE改为16也没有报错。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

这样子就可以了。