typing

Author: Borda

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 = int(pl_module.flops_per_batch)
+            flops_per_batch = 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/overrides/distributed.py

@@ -163,9 +163,7 @@ def _register_ddp_comm_hook(
 
 def _sync_module_states(module: torch.nn.Module) -> None:
     """Taken from https://github.com/pytorch/pytorch/blob/v2.0.0/torch/nn/parallel/distributed.py#L675-L682."""
-    parameters_to_ignore = (
-        set(module._ddp_params_and_buffers_to_ignore) if hasattr(module, "_ddp_params_and_buffers_to_ignore") else set()
-    )
+    parameters_to_ignore = set(getattr(module, "_ddp_params_and_buffers_to_ignore", []))
     from torch.distributed.distributed_c10d import _get_default_group
     from torch.distributed.utils import _sync_module_states as torch_sync_module_states