code: def test_proxyless_bp_hook_once():
      class DistributedModule(LightningModule):
          def __init__(self):
              super().__init__()
              self.op = ProxylessMixedLayer(
                  {
                      "linear0": nn.Linear(3, 3, bias=False),
                      "linear1": ProxylessMixedLayer(
                          {
                              "linear10": nn.Linear(3, 3, bias=False),
                              "linear11": nn.Linear(3, 3, bias=False),
                          },
                          nn.Parameter(torch.randn(2)),
                          nn.Softmax(-1),
                          "linear1",
                      ),
                  },
                  nn.Parameter(torch.randn(2)),
                  nn.Softmax(-1),
                  "linear",
              )
  
          def configure_optimizers(self):
              return torch.optim.SGD(self.op.arch_parameters(), 1e-3)
  
          def training_step(self, batch):
              (x,) = batch
              self.op.resample({})
              self.op["linear1"].resample({})
              y = self.op(x).sum()
              return y
  
      dataset = TensorDataset(torch.randn(20, 3))
      dataloader = DataLoader(dataset, batch_size=2, shuffle=True)
  
      trainer = Trainer(
          max_epochs=1,
          accelerator="cpu",
          devices=1,
          num_nodes=1,
          strategy="ddp_find_unused_parameters_true",
          use_distributed_sampler=False,
      )
  
      trainer.fit(DistributedModule(), dataloader)
  
imports: import random
  import pytest
  import torch
  import torch.nn
  from torch.utils.data import TensorDataset, DataLoader
  import nni
  from nni.mutable import frozen
  from nni.nas.hub.pytorch.nasbench201 import OPS_WITH_STRIDE
  from nni.nas.nn.pytorch import Repeat
  from nni.nas.oneshot.pytorch.supermodule.base import BaseSuperNetModule
  from nni.nas.oneshot.pytorch.supermodule.proxyless import ProxylessMixedLayer, ProxylessMixedInput, ProxylessMixedRepeat, _iter_tensors
  from pytorch_lightning import LightningModule, Trainer
label_desc: Use image node_id 3 for calling a global function with example usage: test_proxyless_bp_hook_once() and returns: torch, y