はじめに
Pytorchのnn.Moduleはto()とかcuda()でテンソルを違うGPU番号へ移すことができますが、すべてのテンソルが移ってくれるわけではありません。
nn.Moduleの実装のすべてを移さなくてはいけないのです。これはだるい。
解決方法
方法は至ってシンプルです。
実行するときに、環境変数を以下のよう変えて、1つのGPUしか見えなくしてあげましょう。
プログラムでcuda:0と書くと、実際には2番目のGPUに繋いでくれます。
$ CUDA_VISIBLE_DEVICES=2 python main.py
