code: def setup(self, stage: str) -> None:
      self.layer = torch.nn.Linear(1, 1)
      weights = self.layer.weight.item(), self.layer.bias.item()
      self.rank_0_weights = self.trainer.strategy.broadcast(weights)
  
file_path: pytorch-lightning/tests/tests_pytorch/overrides/test_distributed.py
imports: from typing import Iterable
  import pytest
  import torch
  from lightning.fabric.utilities.data import has_len
  from lightning.pytorch import LightningModule, Trainer, seed_everything
  from lightning.pytorch.overrides.distributed import UnrepeatedDistributedSampler, _IndexBatchSamplerWrapper
  from torch.utils.data import BatchSampler, SequentialSampler
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 1 for calling the MyModel obj's underlying member method code with example usage: obj.setup(stage) without return types