colab 中 gpu 可用

import torch
from torch import nn, optim
from torchtext import data, datasets
 
if torch.cuda.is_available():
  print('PyTorch can use GPUs!')
else:
  print('PyTorch cannot use GPUs.') 

>>PyTorch can use GPUs!

但是在执行脚本

# ! 表示在 colab jupyter notebook 中执行
!python generate.py -gid 2

提示 gpu 不可用
Warning: No GPU found! Using the CPU instead. The iterations will be slow.

原因：
colab 中只有一块 GPU，但是脚本设置 gid 为 2，表示GPU有三块。

generate.py 中也有一段代码，表示GPU有三块

vq_parser.add_argument("-gid", "--gpu_id", type=str, default=2, dest="gpu_id")

解决方法：
执行脚本改为：

!python generate.py -gid 0

vq_parser.add_argument("-gid", "--gpu_id", type=str, default=0, dest="gpu_id")