code: def test_step(self, batch, batch_idx):
      current = self.layer.weight.item(), self.layer.bias.item()
      assert self.rank_0_weights == current
      gathered = self.all_gather(current)
      # the weights have been synced
      assert all(torch.all(t == t[0]) for t in gathered), gathered
  
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 2 for calling the MyModel obj's underlying member method code with example usage: obj.test_step(batch, batch_idx) without return types