Resuming uses more VRAM than starting from scratch with same parameters

I’m able to run with the following parameters

MODEL_DIM = 512
TEXT_SEQ_LEN = 256
DEPTH = 64
HEADS = 62
DIM_HEAD = 64
BATCH_SIZE=16

using the VQGanVAE1024 vae, however when I try to resume with these parameters, I am hit with cuda out of memory errors. I can’t get past one step until I lower the batch size to 10. Any ideas why resuming with the same parameters would be more memory intensive?

[awilson9]

The fix looks to be mapping the call to torch.load to the cpu rather than the gpu

loaded_obj = torch.load(str(dalle_path), map_location='cpu')

Adding this map location argument got rid of the memory errors for me

Edit: from pytorch docs

When you call torch.load() on a file which contains GPU tensors, those tensors will be loaded to GPU by default. You can call torch.load(…, map_location=‘cpu’) and then load_state_dict() to avoid GPU RAM surge when loading a model checkpoint.

[afiaka87]

@lucidrains here ya go 😃 https://github.com/lucidrains/DALLE-pytorch/pull/118

thanks @awilson9