code: def all_gather_ddp_spawn_fn(strategy):
      rank = strategy.local_rank
      world_size = strategy.num_processes
  
      tensor1 = torch.ones(8, requires_grad=True)
      tensor2 = torch.ones((8, 16, 32), requires_grad=True)
  
      tensor1_gathered = strategy.all_gather(tensor1, sync_grads=True)
      tensor2_gathered = strategy.all_gather(tensor2, sync_grads=True)
  
      tensor1_gathered = tensor1_gathered * rank
      tensor2_gathered = tensor2_gathered * rank
  
      tensor1_gathered.sum().backward()
      tensor2_gathered.sum().backward()
  
      grad1 = torch.zeros_like(tensor1.grad).fill_(
          torch.arange(world_size).sum().float()
      )
      grad2 = torch.zeros_like(tensor2.grad).fill_(
          torch.arange(world_size).sum().float()
      )
  
      assert torch.allclose(grad1, tensor1.grad)
      assert torch.allclose(grad2, tensor2.grad)
  
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 1 for calling a global function with example usage: all_gather_ddp_spawn_fn(strategy) without return types