Skip to content

Instantly share code, notes, and snippets.

@ij96

ij96/torch_get_number_of_weights.py

Created March 2, 2021 13:10
Show Gist options
  • Save ij96/854f63fe5fba593a1db0f5bf422031a8 to your computer and use it in GitHub Desktop.
Save ij96/854f63fe5fba593a1db0f5bf422031a8 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
torch_get_number_of_weights.py
def get_number_of_weights(model, trainable_only=True):
if trainable_only:
return sum(p.numel() for p in model.parameters() if p.requires_grad)
return sum(p.numel() for p in model.parameters())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment