add testing

Author: SkafteNicki

tests/tests_pytorch/utilities/test_model_summary.py

@@ -324,19 +324,33 @@ def test_empty_model_size(max_depth):
         pytest.param("mps", marks=RunIf(mps=True)),
     ],
 )
-def test_model_size_precision(tmp_path, accelerator):
-    """Test model size for half and full precision."""
-    model = PreCalculatedModel()
+@pytest.mark.parametrize("precision", ["16-true", "32-true", "64-true"])
+def test_model_size_precision(tmp_path, accelerator, precision):
+    """Test model size for different precision types."""
+    model = PreCalculatedModel(precision=int(precision.split("-")[0]))
 
     # fit model
     trainer = Trainer(
-        default_root_dir=tmp_path, accelerator=accelerator, devices=1, max_steps=1, max_epochs=1, precision=32
+        default_root_dir=tmp_path, accelerator=accelerator, devices=1, max_steps=1, max_epochs=1, precision=precision
     )
     trainer.fit(model)
     summary = summarize(model)
     assert model.pre_calculated_model_size == summary.model_size
 
 
+def test_model_size_warning_on_unsupported_precision():
+    """Test that a warning is raised when the precision is not supported."""
+    model = PreCalculatedModel(precision=32)  # fallback to 32 bits
+
+    # supported precision by lightning but not by the model summary
+    trainer = Trainer(max_epochs=1, precision="16-mixed")
+    trainer.fit(model)
+
+    with pytest.warns(UserWarning, match="Precision 16-mixed is not supported by the model summary.*"):
+        summary = summarize(model)
+        assert model.pre_calculated_model_size == summary.model_size
+
+
 def test_lazy_model_summary():
     """Test that the model summary can work with lazy layers."""
     lazy_model = LazyModel()