torch.cuda.set_divice()

我申请了两块GPU,然后看两张显卡的编号

import torch
torch.cuda.set_device(0)                            # 设置当前cuda设备编号为1
print("当前cuda设备是", torch.cuda.current_device()) # 获取当前cuda设备

import torch
torch.cuda.set_device(1)                            # 设置当前cuda设备编号为1
print("当前cuda设备是", torch.cuda.current_device()) # 获取当前cuda设备

import torch
torch.cuda.set_device(2)                            # 设置当前cuda设备编号为1
print("当前cuda设备是", torch.cuda.current_device()) # 获取当前cuda设备