在程序最开始：

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'  # 使用多卡
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

使用单卡（前面也只写一个GPU），将数据和模型放到GPU上：

input = data.to(device)
model = MyModule(xxx).to(device)

使用多卡，针对模型，进行并行化处理：

model = MyModule(xxx)
if torch.cuda.device_count() > 1:
	print("Let's use", torch.cuda.device_count(), "GPUs!")
	model = nn.DataParallel(model)
model = model.to(device)

参考：

1. RuntimeError: CUDA out of memory 解决办法：https://blog.csdn.net/Wadewhl/article/details/123891113
2. 【NLP】pytorch中CPU、GPU的使用（仅CPU、单机多卡、多机多卡）：https://blog.csdn.net/sunflower_sara/article/details/109675174
3. Pytorch的nn.DataParallel：https://zhuanlan.zhihu.com/p/102697821