code: def test_all_gather_sync_grads(tmpdir):
      class TestModel(BoringModel):
          training_step_called = False
  
          def training_step(self, batch, batch_idx):
              self.training_step_called = True
              tensor = torch.rand(
                  2, 2, requires_grad=True, device=self.device
              )
              gathered_tensor = self.all_gather(tensor, sync_grads=True)
              assert gathered_tensor.shape == torch.Size([2, 2, 2])
              return gathered_tensor.sum()
  
      model = TestModel()
      trainer = Trainer(
          default_root_dir=tmpdir,
          limit_train_batches=1,
          limit_val_batches=0,
          max_epochs=1,
          accelerator="gpu",
          devices=2,
          strategy="ddp",
          enable_progress_bar=False,
          enable_model_summary=False,
          enable_checkpointing=False,
      )
      trainer.fit(model)
      assert model.training_step_called
  
file_path: pytorch-lightning/tests/tests_pytorch/utilities/test_all_gather_grad.py
imports: import numpy
  import pytest
  import torch
  from lightning.pytorch import Trainer
  from lightning.pytorch.demos.boring_classes import BoringModel
  from tests_pytorch.core.test_results import spawn_launch
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 4 for calling a global function with example usage: test_all_gather_sync_grads(tmpdir) and returns: gathered_tensor