Skip to content

Instantly share code, notes, and snippets.

@yatakeke

yatakeke/pytorch_model_save_and_load.py

Last active November 29, 2019 05:49
Show Gist options
  • Save yatakeke/edbe28a568c4d2b28b4a36b21351c032 to your computer and use it in GitHub Desktop.
Save yatakeke/edbe28a568c4d2b28b4a36b21351c032 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
pytorch_model_save_and_load.py
import torch
# save parameters
torch.save(model.state_dict(), PATH)
# load parameters
model = ModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment