pytorch版本0.3.1.post2的设置设备

时间:2019-05-07 16:06:32

标签: python python-3.x pytorch torch

在较新版本的pytorch中,可以使用torch.device设置设备。对于某些旧代码,我需要使用割炬版本0.3.1.post2。如何为该版本设置设备?

1 个答案:

答案 0 :(得分:0)

据我所知,您可以使用set_device函数。但这是不鼓励的。请参阅reference

建议的方法是,只需设置CUDA_VISIBLE_DEVICES环境变量。您可以按以下方式运行脚本。

CUDA_VISIBLE_DEVICES=GPU_ID python script_name.py

在您的程序中,您只需使用.cuda()即可使用GPU。 (例如model=model.cuda()