pytorch指定使用的gpu设备

最近在使用pytorch写程序，想在指定的GPU设备上执行。

假设函数func( )使用到了GPU设备：

则在func()函数内部或者调用func()的地方加上:

os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

我之前尝试在此处添加torch.cuda.set_device(self.gpu_id)，没作用，记录下。