def METHOD_NAME(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()