From 16299b805a05bbe27f79c1fc12424624e29a8bfe Mon Sep 17 00:00:00 2001
From: Rio H <hoshrio@gmail.com>
Date: Tue, 27 Jul 2021 21:28:12 -0700
Subject: [PATCH] Lint Files and fix test_memory.py

---
 pytorch_lightning/core/lightning.py   | 4 ++--
 pytorch_lightning/utilities/memory.py | 2 +-
 tests/callbacks/test_quantization.py  | 2 +-
 tests/utilities/test_memory.py        | 8 ++++----
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index e7b254e473bcaf..136bdc7c1a5b91 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -1963,8 +1963,8 @@ def model_size(self) -> float:
             Will be removed in v1.7.0
         """
         rank_zero_deprecation(
-            'The `LightningModule.model_size` property was deprecated in v1.5 and will be removed in v1.7. Please '
-            'use the get_model_size_mb method under utilities/memory.py'
+            "The `LightningModule.model_size` property was deprecated in v1.5 and will be removed in v1.7. Please "
+            "use the get_model_size_mb method under utilities/model_helpers.py"
         )
 
         tmp_name = f"{uuid.uuid4().hex}.pt"
diff --git a/pytorch_lightning/utilities/memory.py b/pytorch_lightning/utilities/memory.py
index f50d965711caac..ab8fae40213071 100644
--- a/pytorch_lightning/utilities/memory.py
+++ b/pytorch_lightning/utilities/memory.py
@@ -13,8 +13,8 @@
 # limitations under the License.
 
 import gc
-import uuid
 import os
+import uuid
 
 import torch
 from torch.nn import Module
diff --git a/tests/callbacks/test_quantization.py b/tests/callbacks/test_quantization.py
index 3867e962b740cb..5d39cff53d33d3 100644
--- a/tests/callbacks/test_quantization.py
+++ b/tests/callbacks/test_quantization.py
@@ -20,8 +20,8 @@
 from pytorch_lightning import seed_everything, Trainer
 from pytorch_lightning.callbacks import QuantizationAwareTraining
 from pytorch_lightning.metrics.functional.mean_relative_error import mean_relative_error
-from pytorch_lightning.utilities.memory import get_model_size_mb
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.memory import get_model_size_mb
 from tests.helpers.datamodules import RegressDataModule
 from tests.helpers.runif import RunIf
 from tests.helpers.simple_models import RegressionModel
diff --git a/tests/utilities/test_memory.py b/tests/utilities/test_memory.py
index 58c3538864cb84..70cf008d81578c 100644
--- a/tests/utilities/test_memory.py
+++ b/tests/utilities/test_memory.py
@@ -11,11 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import math
 import torch
+import torch.nn as nn
 
 from pytorch_lightning.utilities.memory import get_model_size_mb, recursive_detach
 from tests.helpers import BoringModel
-import torch.nn as nn
 
 
 def test_recursive_detach():
@@ -38,8 +39,7 @@ def test_get_model_size_mb():
     size_bytes = get_model_size_mb(model)
 
     # Size will be python version dependent.
-    assert size_bytes == 0.001319 or size_bytes == 0.001409
-
+    assert math.isclose(size_bytes, 0.001319, rel_tol=0.1)
 
 def test_get_sparse_model_size_mb():
     class BoringSparseModel(BoringModel):
@@ -50,4 +50,4 @@ def __init__(self):
     model = BoringSparseModel()
     size_bytes = get_model_size_mb(model)
 
-    assert size_bytes == 0.001511 or size_bytes == 0.001665
+    assert math.isclose(size_bytes, 0.001511, rel_tol=0.1)