typing

Author: Borda

src/lightning/fabric/strategies/xla_fsdp.py

@@ -295,6 +295,7 @@ def clip_gradients_norm(
     ) -> Tensor:
         """Clip gradients by norm."""
         self.precision.unscale_gradients(optimizer)
+        assert callable(module.clip_grad_norm_)
         return module.clip_grad_norm_(max_norm=max_norm, norm_type=norm_type)
 
     @override

src/lightning/fabric/utilities/init.py

@@ -67,7 +67,8 @@ def _materialize(module: Module, device: _DEVICE) -> None:
             f"Materialization requires that the `{type(module).__name__}.reset_parameters` method is implemented."
             " This method is used to initialize any children parameters or buffers in this module."
         )
-    module.reset_parameters()
+    if callable(module.reset_parameters):
+        module.reset_parameters()
 
 
 def _materialize_meta_tensors(module: Module, device: _DEVICE) -> None:

src/lightning/pytorch/callbacks/finetuning.py

@@ -133,7 +133,7 @@ def flatten_modules(modules: Union[Module, Iterable[Union[Module, Iterable]]]) -
 
         if isinstance(modules, Iterable):
             _flatten_modules = []
-            for m in modules:  # type: ignore[union-attr]
+            for m in modules:
                 _flatten_modules.extend(BaseFinetuning.flatten_modules(m))
 
             _modules = iter(_flatten_modules)

src/lightning/pytorch/callbacks/throughput_monitor.py

@@ -125,7 +125,7 @@ def _update(self, trainer: "Trainer", pl_module: "LightningModule", batch: Any,
             self._lengths[stage] += self.length_fn(batch)
 
         if hasattr(pl_module, "flops_per_batch"):
-            flops_per_batch = pl_module.flops_per_batch
+            flops_per_batch = int(pl_module.flops_per_batch)
         else:
             rank_zero_warn(
                 "When using the `ThroughputMonitor`, you need to define a `flops_per_batch` attribute or property"

src/lightning/pytorch/plugins/precision/double.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from collections.abc import Generator
 from contextlib import AbstractContextManager, contextmanager
-from typing import Any, Literal
+from typing import Any, Literal, Iterable
 
 import torch
 import torch.nn as nn
@@ -71,6 +71,7 @@ class LightningDoublePrecisionModule(_DeviceDtypeModuleMixin, nn.Module):
         pl_module: the model to wrap
 
     """
+    _ddp_params_and_buffers_to_ignore: Iterable[str]
 
     def __init__(self, pl_module: "pl.LightningModule") -> None:
         super().__init__()