def average_shard_parameters(self) -> int:
        """Returns the number of parameters in this module."""

        def partitioned_size(p: Parameter) -> int:
            return p.partitioned_size() if RequirementCache("deepspeed<0.6.6") else p.partition_numel()

        return sum(partitioned_size(p) if not _tensor_has_shape(p) else 0 for p in self._module.parameters())
Returns the number of parameters in this module.
src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py