Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

protect ByteCounter (#10300)

This commit is contained in:
Rohit Gupta 2021-11-02 14:05:40 +05:30 committed by GitHub
parent 3cd65b592b
commit f44b8a7aab
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
pytorch_lightning/utilities/memory.py
View File

@ -24,7 +24,7 @@ from torch.nn import Module
from pytorch_lightning.utilities.apply_func import apply_to_collection
class ByteCounter:
class _ByteCounter:
"""Accumulate and stores the total bytes of an object."""
def __init__(self) -> None:
@ -184,7 +184,7 @@ def get_model_size_mb(model: Module) -> float:
Returns:
Number of megabytes in the parameters of the input module.
"""
model_size = ByteCounter()
model_size = _ByteCounter()
torch.save(model.state_dict(), model_size)
size_mb = model_size.nbytes / 1e6
return size_mb