Hi ptrblck,
Thanks for the suggestion, unfortunetly I don’t think it’s that as my training loop is this:
for epoch in range(num_epochs):
    torch.cuda.empty_cache()
    train_one_epoch(model, optimizer, data_loader_train, device, epoch, print_freq=1)
    lr_scheduler.step()
    print('Epoch done - Beginning evalutation')
    torch.cuda.empty_cache()
    evaluate(model, data_loader_test, device=torch.device('cpu'))
    torch.cuda.empty_cache()

With the empty cache statements to try to fix this problem, and the train and evaluate functions available here.