多个GPU时如何设置选择哪个GPU来运行

[xiaofeihhu]

服务器目前有2个GPU，每个16G显存。在运行时好像每次默认只会选择第1个GPU，另外一个空闲着，并且会出现下面的报错。
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 14.57 GiB total capacity; 2.53 GiB already allocated; 12.75 MiB free; 2.60 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

[nishuone028]

推理代码哪里 去找cuda设置你的卡