小白学GAN系列3——torch.device

2022-10-11 08:26:49
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

  这段代码一般写在读取数据之前,torch.device代表将torch.Tensor分配到的设备的对象。torch.device包含一个设备类型(‘cpu’或‘cuda’)和可选的设备序号。如果设备序号不存在,则为当前设备。如:torch.Tensor用设备构建‘cuda’的结果等同于‘cuda:X’,其中X是torch.cuda.current_device()的结果。

mytensor = my_tensor.to(device)

  这行代码的意思是将所有最开始读取数据时的tensor变量copy一份到device所指定的GPU上去,之后的运算都在GPU上进行。

  这句话需要写的次数等于需要保存GPU上的tensor变量的个数;一般情况下这些tensor变量都是最开始读数据时的tensor变量,后面衍生的变量自然也都在GPU上

  • 作者:崔可爱的胖肉肉
  • 原文链接:https://blog.csdn.net/o960604/article/details/109486153
    更新时间:2022-10-11 08:26:49