Using the option --gpu=cuda:0 instead of --gpu=0Or update the code torch.cuda.set_device(args.gpu) with torch.cuda.set_device(int(args.gpu))