Pytorch GPU memory increase after load operation
I have a pytorch model which is of the size 386MB, but when i load the model

state = torch.load(f, flair.device)


My GPU memory takes up to 900MB, why does this happen and is there a way to resolve this?

This is how i save the model

model_state = self._get_state_dict()

# additional fields for model checkpointing
model_state["optimizer_state_dict"] = optimizer_state
model_state["scheduler_state_dict"] = scheduler_state
model_state["epoch"] = epoch
model_state["loss"] = loss

torch.save(model_state, str(model_file), pickle_protocol=4)


It's probably the optimizer_state that takes the extra space. Some optimizers (e.g., Adam) track statistics of each trainable parameter, such as first and second order moments. This information takes up space, as you can tell.

You can load to CPU first:

state = torch.load(f, map_location=torch.device('cpu'))

