From d58c28cb75dfb8789390bdef218854e1c8ce1a2e Mon Sep 17 00:00:00 2001 From: awaelchli Date: Fri, 8 Mar 2024 12:44:18 +0100 Subject: [PATCH] Update warning about `UninitializedParameter` in `ModelSummary` (#19597) update warning --- .../pytorch/utilities/model_summary/model_summary.py | 5 +++-- tests/tests_pytorch/utilities/test_model_summary.py | 6 +----- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/src/lightning/pytorch/utilities/model_summary/model_summary.py b/src/lightning/pytorch/utilities/model_summary/model_summary.py index 1d5612736a2ab..ef2827d3b7eed 100644 --- a/src/lightning/pytorch/utilities/model_summary/model_summary.py +++ b/src/lightning/pytorch/utilities/model_summary/model_summary.py @@ -463,8 +463,9 @@ def _is_lazy_weight_tensor(p: Tensor) -> bool: if isinstance(p, UninitializedParameter): warning_cache.warn( - "A layer with UninitializedParameter was found. " - "Thus, the total number of parameters detected may be inaccurate." + "The total number of parameters detected may be inaccurate because the model contains" + " an instance of `UninitializedParameter`. To get an accurate number, set `self.example_input_array`" + " in your LightningModule." ) return True return False diff --git a/tests/tests_pytorch/utilities/test_model_summary.py b/tests/tests_pytorch/utilities/test_model_summary.py index c493c5043fa1a..290dfb67faf7d 100644 --- a/tests/tests_pytorch/utilities/test_model_summary.py +++ b/tests/tests_pytorch/utilities/test_model_summary.py @@ -345,11 +345,7 @@ def test_lazy_model_summary(): lazy_model = LazyModel() summary = ModelSummary(lazy_model) - with pytest.warns( - UserWarning, - match=r"A layer with UninitializedParameter was found. " - r"Thus, the total number of parameters detected may be inaccurate.", - ): + with pytest.warns(UserWarning, match="The total number of parameters detected may be inaccurate."): assert summary.total_parameters == 0 assert summary.trainable_parameters == 0