本文適合多GPU的機器，並且每個使用者需要單獨使用GPU訓練。

雖然pytorch提供了指定gpu的幾種方式，但是使用不當的話會遇到out of memory的問題，主要是因為pytorch會在第0塊gpu上初始化，並且會佔用一定空間的視訊記憶體。這種情況下，經常會出現指定的gpu明明是空閒的，但是因為第0塊gpu被佔滿而無法執行，一直報out of memory錯誤。解決方案如下：

1. 指定環境變數，遮蔽第0塊gpu

CUDA_VISIBLE_DEVICES = 1 main.py

這句話表示只有第1塊gpu可見，其他gpu不可用，此時要注意第1塊gpu已經變成第0塊，因此程式碼裡應該使用編號0來指定gpu。如果依然使用cuda:1

import os

#多塊使用逗號隔開
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

2. 使用torch.cuda.device

with torch.cuda.device(1):
	...

注意這種情況下gpu編號也會改變

關於這個問題的討論，感興趣的也可以去這裡檢視cuda out of memory error when GPU0 memory is fully utilized