SimonW:

k is still on cpu. Add NN = NN.cuda().


Sorrry to hijack this post, but is there a way to use condition to automatically switch between gpu and cpu depending on the hardware available ?
Sometimes I am working on my macbook, and I am trying to find a way not to modify this part every time.
I found it is possible to use parsing to detect GPU availability, but I don’t know how to do that on network.
http://pytorch.org/docs/master/notes/cuda.html 10
import argparse
import torch

parser = argparse.ArgumentParser(description='PyTorch Example')
parser.add_argument('--disable-cuda', action='store_true', help='Disable CUDA')
args = parser.parse_args()
args.cuda = not args.disable_cuda and torch.cuda.is_available()

if args.cuda:
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
        # something like torch.network.cuda() ??   

##Your model below