Taking up memory on the primary GPU loading from checkpoint

[pc4653]

Hi,
I'm wondering what is taking up memory on the main GPU when I am resuming training from a checkpoint. As a result i cannot train with a larger patch size, which led to memory limitation.

I think this is caused by the distributed training system and wonder if there is any way to either avoid the memory cost or distribute it evenly to other GPUs?
Many thanks!

[rosinality]

I think it is related to #31. You can apply PR or try to use torch.load(args.ckpt, map_location=lambda storage, loc: storage).

[pc4653]

you are correct, thanks!