Not at all. There are very few lines of code where I access CUDA.
The first part is where I move the model to CUDA
net = Net()
net = torch.nn.DataParallel(net)
net.to(device)

What I showed in the code snippet at the top is sort of the second part, where I get the minibatch, do forward pass and back propagation etc. What I didn’t show, though, is a just couple of lines to take snapshot of the model and also print the loss directly to a file, as below.
# save the model at specified steps
if step % snapshot == 0 and step >= snapshot:
    file_name = 'ho-rcnn-model-step' + step
    torch.save(net.state_dict(), os.path.join(path, file_name))

# print statistics
if step % printstats == 0:
    with open(path_stats, 'ab') as fid:
        fid.write('\n' + str(datetime.datetime.now()))
        fid.write('\nStep: {}, loss of current batch is {}\n'.format(step, loss))

Apart from that, there isn’t anymore usage of CUDA, or reference to any variables stored in CUDA, which is why I think the problem might be caused within some of the encapsulated torch modules