pytorch多GPU的正确使用方式

假如你有多个GPU（>2），此处假设你有8块，并且前4块已经使用了，那么怎么才能从第5块开始用呢？
常见错误：

    raise AssertionError("Invalid device id")
AssertionError: Invalid device id

THCudaCheck FAIL file=..\torch\csrc\cuda\Module.cpp line=59 error=101 : invalid
device ordinal
Traceback (most recent call last):
  File "train2.py", line 26, in <module>
    torch.cuda.set_device(4)
  File "C:\ProgramData\Anaconda3\lib\site-packages\torch\cuda\__init__.py", line
 281, in set_device
    torch._C._cuda_setDevice(device)
RuntimeError: cuda runtime error (101) : invalid device ordinal at ..\torch\csrc
\cuda\Module.cpp:59

解决方法：
我们知道GPU的编号是从0开始的，这8块的编号就是0，1，2，3，4，5，6，7
现在要从4开始使用，只用4和5
1.在你的训练脚本import 结束的时候放下面两句：

torch.cuda.set_device(4)  #指定主显卡，不指定不行，默认是0，但是0被占用了
os.environ["CUDA_VISIBLE_DEVICES"] = "4,5"

2.脚本里面需要放：

device = torch.device("cuda:4" if torch.cuda.is_available() else "cpu")

3.还需要改：

self.net = torch.nn.DataParallel(self.net, device_ids=[4,5])