只需要在用到device的地方,如
torch.Tensor([(self._dim + 1) / 2], device=self.device)
加上这句话即可,表示又cuda就用cuda,没有的话就用cpu
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ps:可以把"cuda:0"换成"cuda:1"试试
只需要在用到device的地方,如
torch.Tensor([(self._dim + 1) / 2], device=self.device)
加上这句话即可,表示又cuda就用cuda,没有的话就用cpu
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ps:可以把"cuda:0"换成"cuda:1"试试