Balance memory for DataParallel

[dyukha (Dmitrii Avdiukhin)]

Thanks for the reply! I have to use DataParallel because of issues with DDP: Distributed Data Parallel example - "process 0 terminated with exit code 1" - #3 by dyukha