From f44b8a7aaba9fb33ca3ac86db1f6489808c4d5a7 Mon Sep 17 00:00:00 2001
From: Rohit Gupta <rohitgr1998@gmail.com>
Date: Tue, 2 Nov 2021 14:05:40 +0530
Subject: [PATCH] protect ByteCounter (#10300)

---
 pytorch_lightning/utilities/memory.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pytorch_lightning/utilities/memory.py b/pytorch_lightning/utilities/memory.py
index e810c70b3429b..b80a2b94bd96b 100644
--- 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