pytorch使用任意编号的多块gpu训练

查了一些资料，要么把全部gpu都用上，要不只能用一块，假如我4块gpu，前面2块别人在用，我想把后面2块全用上，那就不行了。所以，本文是为了解决这个需求的，相关操作已经在注释中。

# 测试使用指定的gpu编号
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

import os
os.environ["CUDA_