Hello, I am not entirely sure about this, but when you execute:
if use_cuda:
   model.cuda()
   model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
   cudnn.benchmark = True

you are storing the entire model on GPU 0 (defaul cuda device) but then  you are (when using DataParalell) storing a copy of the model on GPU 0 and GPU 1. Maybe that is why your GPU 0 is running out of memory so quickly.  You should remove the model.cuda() line. Let me know if this helps.
Diego