拿到了师兄的程序,但是自己手头的电脑无法使用cpu训练,虽然也可以使用colab在线运行,但是还是本地的会舒服一些,anyway,就得将代码稍作一下修改
方法:
在训练之前加入这一句
#选择使用gpu,如果没有gpu就使用cpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
然后,在训练的代码中会找到 .cuda() ,将所有的都 替换成为.to(device) 即可
(在网络后边和输入输出后边都需要改)
即将