Thank you for the answer.
In terms of assigning to the gpu I have a check after the creation of the model to send it to the device
    cuda = torch.cuda.is_available()
    if cuda and gpu_to_use > -1:
        use_cuda = True
        torch.cuda.device(gpu_to_use)
        torch.cuda.seed()
        device = torch.device("cuda", gpu_to_use)
        print("Using CUDA.")
    else:
        use_cuda = False
        device = torch.device("cpu")
        print("Using CPU.")  
  net = model.to(device)

Once I’ve done this, in the training and testing function, I give the device and have basically something like this (in the iteration over batches):
out = net(sample_batched['tensor'].to(device).view(batch_size, -1, max_size * 20))
target = sample_batched['interaction'].to(device)
loss = criterion(out, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()

I forgot to say that I am currently working on pytorch 0.4.0,  I don’t know if that is of any help.