code: def training_step(self, batch, batch_idx):
      # emulate gans training
      opt_gen, opt_dis = self.optimizers()
  
      # Note: Be careful, don't log on the same key in self.log in both closure
      # as they will be aggregated together on epoch_end
  
      world_size = torch_distrib.get_world_size(
          torch_distrib.group.WORLD
      )
      assert world_size == 2
  
      make_gen_optimizer_step = batch_idx % 2 == 1
      make_dis_optimizer_step = batch_idx % 4 == 0
  
      def compute_loss():
          x = batch[0]
          x = F.dropout(x, 0.1)
          predictions = self(x)
          predictions = F.dropout(predictions, 0.1)
          loss_ones = self.loss_ones(None, predictions)
          loss_zeros = self.loss_zeros(None, predictions)
          return loss_ones, loss_zeros
  
      def make_manual_backward(
          loss, retain_graph=False, make_optimizer_step=True
      ):
          self.manual_backward(loss, retain_graph=retain_graph)
          if make_optimizer_step:
              grad_clone = self.layer.weight.grad.clone()
              assert self.manual_sync_grad()
              self.layer.weight.grad /= world_size
              assert torch.equal(self.layer.weight.grad, grad_clone)
  
      def gen_closure():
          loss_ones_gen, _ = compute_loss()
          make_manual_backward(
              loss_ones_gen,
              retain_graph=True,
              make_optimizer_step=make_gen_optimizer_step,
          )
          make_manual_backward(
              loss_ones_gen, make_optimizer_step=make_gen_optimizer_step
          )
  
      def dis_closure():
          loss_ones_gen, _ = compute_loss()
          make_manual_backward(
              loss_ones_gen,
              retain_graph=True,
              make_optimizer_step=make_dis_optimizer_step,
          )
          make_manual_backward(
              loss_ones_gen, make_optimizer_step=make_dis_optimizer_step
          )
  
      # this will accumulate gradients for 2 batches and then call opt_gen.step()
      with opt_gen.toggle_model(sync_grad=make_gen_optimizer_step):
          gen_closure()
          if make_gen_optimizer_step:
              opt_gen.step()
              opt_gen.zero_grad()
  
      with opt_dis.toggle_model(sync_grad=make_dis_optimizer_step):
          dis_closure()
          if make_dis_optimizer_step:
              opt_dis.step()
              opt_dis.zero_grad()
  
imports: import collections
  import contextlib
  from copy import deepcopy
  from unittest import mock
  from unittest.mock import ANY, call, patch
  import pytest
  import torch
  import torch.distributed
  import torch.nn.functional
  from lightning.fabric.utilities.exceptions import MisconfigurationException
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
  from lightning.pytorch import Trainer, seed_everything
  from lightning.pytorch.demos.boring_classes import BoringModel, ManualOptimBoringModel
  from lightning.pytorch.strategies import Strategy
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 1 for calling the TestManualOptimizationDDPModelToggleModel obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: loss_ones,  loss_zeros