Linux运行python文件指定GPU

最近用GPU小集群，指定GPU的代码小笔记：可用装饰函数,deiviceid也可以用int再str()，无GPU自动CPU，运行无指定-d自动cuda0，指定-d使用指定GPU（nvidia-smi查看）

运行代码：python main.py -d 1

main.py内容：

import click

@click.command()
@click.option('--deviceid','-d', help="device =__", type=str, required=True, default="0") 

def main(deviceid):

    #DEVICE
    if torch.cuda.is_available():
        device = "cuda:"+deviceid
        print('using device ' + device) 
    else:
        device = 'cpu'    
        print('using device ' + device)    
    device = torch.device(device)

if __name__ == "__main__":
    main()