Search code examples
pythonpython-3.xpytorchtorch

Setting device for pytorch version 0.3.1.post2

In newer versions of pytorch, you can set the device by using torch.device. I need to use torch version 0.3.1.post2 for some legacy code. How do I set the device for this version?

Solution

  • As far as I know, you can use the set_device function. But this is not encouraged. Please see the reference.

    The suggested method is, you just set the CUDA_VISIBLE_DEVICES environmental variable. You can run your script as follows.

    CUDA_VISIBLE_DEVICES=GPU_ID python script_name.py

    In your program, you can just simply use .cuda() to use GPUs. (e.g., model=model.cuda())