update model message (#20753)

* update model message
* nitpick_ignore

(cherry picked from commit 6cf2581)

Author: Borda

.github/workflows/docs-build.yml

@@ -72,7 +72,7 @@ jobs:
           lfs: ${{ matrix.pkg-name == 'pytorch' }}
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.9"
+          python-version: "3.10"
 
       - name: List notebooks
         if: ${{ matrix.pkg-name == 'pytorch' }}

docs/source-fabric/conf.py

@@ -292,6 +292,8 @@
     ("py:.*", "torch_xla.*"),
     ("py:class", "transformer_engine.*"),
     ("py:class", "bitsandbytes.*"),
+    # loggers
+    ('py:class', 'tensorboardX.SummaryWriter'),  # todo: this is unexpected as the imports locally works
 ]
 
 # -- Options for todo extension ----------------------------------------------

docs/source-pytorch/conf.py

@@ -377,6 +377,7 @@ def _load_py_module(name: str, location: str) -> ModuleType:
     # missing in generated API
     ("py:exc", "MisconfigurationException"),
     # TODO: generated list of all existing ATM, need to be fixed
+    ('py:class', 'tensorboardX.SummaryWriter'),
     ("py:class", "AveragedModel"),
     ("py:class", "CometExperiment"),
     ("py:meth", "DataModule.__init__"),

src/lightning/pytorch/trainer/connectors/callback_connector.py

@@ -18,7 +18,7 @@
 from datetime import timedelta
 from typing import Optional, Union
 
-from lightning_utilities import module_available
+from lightning_utilities.core.imports import RequirementCache
 
 import lightning.pytorch as pl
 from lightning.fabric.utilities.registry import _load_external_callbacks
@@ -93,7 +93,7 @@ def _configure_checkpoint_callbacks(self, enable_checkpointing: bool) -> None:
                     " but found `ModelCheckpoint` in callbacks list."
                 )
         elif enable_checkpointing:
-            if module_available("litmodels") and self.trainer._model_registry:
+            if RequirementCache("litmodels >=0.1.7") and self.trainer._model_registry:
                 trainer_source = inspect.getmodule(self.trainer)
                 if trainer_source is None or not isinstance(trainer_source.__package__, str):
                     raise RuntimeError("Unable to determine the source of the trainer.")
@@ -103,12 +103,11 @@ def _configure_checkpoint_callbacks(self, enable_checkpointing: bool) -> None:
                 else:
                     from litmodels.integrations.checkpoints import LightningModelCheckpoint as LitModelCheckpoint
 
-                model_checkpoint = LitModelCheckpoint(model_name=self.trainer._model_registry)
+                model_checkpoint = LitModelCheckpoint(model_registry=self.trainer._model_registry)
             else:
                 rank_zero_info(
-                    "You are using the default ModelCheckpoint callback."
-                    " Install `litmodels` package to use the `LitModelCheckpoint` instead"
-                    " for seamless uploading to the Lightning model registry."
+                    "Using default `ModelCheckpoint`. Consider installing `litmodels` package to enable"
+                    " `LitModelCheckpoint` for automatic upload to the Lightning model registry."
                 )
                 model_checkpoint = ModelCheckpoint()
             self.trainer.callbacks.append(model_checkpoint)