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)