Lint Files and fix test_memory.py

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"

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

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

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)