pytroch中指定GPU或者CPU寫法

之前習慣用.cpu()，.cuda()來指定．
現在不要顯示的指定是gpu, cpu之類的. 利用.to()來執行

# at beginning of the script
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

...

# then whenever you get a new Tensor or Module
# this won't copy if they are already on the desired device
input = data.to(device)
model = MyModule(...).to(device)

用.to()和.cuda()代碼對比
.cuda()寫法

model = MyRNN()
  if use_cuda:
      model = model.cuda()

  # train
  total_loss = 0
  for input, target in train_loader:
      input, target = Variable(input), Variable(target)
      hidden = Variable(torch.zeros(*h_shape))  # init hidden
      if use_cuda:
          input, target, hidden = input.cuda(), target.cuda(), hidden.cuda()
      ...  # get loss and optimize
      total_loss += loss.data[0]

  # evaluate
  for input, target in test_loader:
      input = Variable(input, volatile=True)
      if use_cuda:
          ...
      ...

現在.to()

# torch.device object used throughout this script
  device = torch.device("cuda" if use_cuda else "cpu")

  model = MyRNN().to(device)

  # train
  total_loss = 0
  for input, target in train_loader:
      input, target = input.to(device), target.to(device)
      hidden = input.new_zeros(*h_shape)  # has the same device & dtype as `input`
      ...  # get loss and optimize
      total_loss += loss.item()           # get Python number from 1-element Tensor

  # evaluate
  with torch.no_grad():                   # operations inside don't track history
      for input, target in test_loader:
          ...

參考：https://www.itency.com/topic/show.do?id=494122