Correct GPU memory allocation

[my3bikaht (Sergey)]

Also take a look at this post:

DataParallel vs DistributedDataParallel

Hi, what is the difference between model = nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu]) and model = nn.DataParallel(model, device_ids=[args.gpu]) ?