code: def test_concat_loader_nested():
      from nni.nas.oneshot.pytorch._dataloader import ConcatLoader
  
      loaders = {
          "a": [
              DataLoader(range(10), batch_size=4),
              DataLoader(range(20), batch_size=6),
          ],
          "b": DataLoader(range(20), batch_size=5),
      }
      dataloader = ConcatLoader(loaders)
      assert len(dataloader) == 11
      for i, (data, batch_index, loader_index) in enumerate(dataloader):
          if i < 3:
              assert len(data) in [2, 4]
              assert loader_index == 0
              assert batch_index == i
          else:
              assert 1 <= loader_index <= 2
  
imports: import math
  from typing import Union
  import pytest
  import torch
  import pytorch_lightning
  from pytorch_lightning import LightningModule, Trainer
  from pytorch_lightning.trainer.states import RunningStage
  from torch.utils.data import DataLoader, Dataset
label_desc: Use image node_id 2 for calling a global function with example usage: test_concat_loader_nested() without return types