code: def test_get_batch_size():
      tensor = torch.rand(2, 3)
      assert get_batch_size(tensor) == 2
      assert get_batch_size([tensor]) == 2
      assert get_batch_size((1, tensor)) == 2
      assert get_batch_size({"tensor": tensor}) == 2
      assert get_batch_size({"dummy": [1], "tensor": tensor}) == 2
      assert get_batch_size({"tensor": [tensor]}) == 2
  
imports: import torch
  from colossalai.pipeline.schedule._utils import get_batch_size, get_micro_batch, merge_batch
label_desc: Use image node_id 1 for calling a global function with example usage: test_get_batch_size() without return types