def setup(self, stage: str) -> None:
        """Called at the beginning of fit (train + validate), validate, test, or predict. This is a good hook when you
        need to build models dynamically or adjust something about them. This hook is called on every process when
        using DDP.

        Args:
            stage: either ``'fit'``, ``'validate'``, ``'test'``, or ``'predict'``

        Example::

            class LitModel(...):
                def __init__(self):
                    self.l1 = None

                def prepare_data(self):
                    download_data()
                    tokenize()

                    # don't do this
                    self.something = else

                def setup(self, stage):
                    data = load_data(...)
                    self.l1 = nn.Linear(28, data.num_classes)

        """
Called at the beginning of fit (train + validate), validate, test, or predict. This is a good hook when you
        need to build models dynamically or adjust something about them. This hook is called on every process when
        using DDP.

        Args:
            stage: either ``'fit'``, ``'validate'``, ``'test'``, or ``'predict'``

        Example::

            class LitModel(...):
                def __init__(self):
                    self.l1 = None

                def prepare_data(self):
                    download_data()
                    tokenize()

                    # don't do this
                    self.something = else

                def setup(self, stage):
                    data = load_data(...)
                    self.l1 = nn.Linear(28, data.num_classes)

        
src/lightning/pytorch/core/hooks.py