GPU的编号一般默认从0开始依次递增,利用代码
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
可以指定你想用的GPU。当你想指定多个的时候可以用
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2,3'
注意,修改指定GPU后
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
此代码块中的0不用改变。
GPU的编号一般默认从0开始依次递增,利用代码
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
可以指定你想用的GPU。当你想指定多个的时候可以用
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2,3'
注意,修改指定GPU后
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
此代码块中的0不用改变。