Update warning about UninitializedParameter in ModelSummary (#19597)

update warning
Lightning-AI · Mar 8, 2024 · d58c28c · d58c28c
1 parent 06eb3cc
commit d58c28c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git 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
@@ -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