protect ByteCounter (#10300)

Lightning-AI · Nov 2, 2021 · f44b8a7 · f44b8a7
1 parent 3cd65b5
commit f44b8a7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pytorch_lightning/utilities/memory.py b/pytorch_lightning/utilities/memory.py
@@ -24,7 +24,7 @@
 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