class_docstr: Test Shunted Transformer
code: def test_base(self):
      dummy_img = torch.rand(1, 3, 16, 16).to(self.device)
  
      model1 = create_model(
          img_size=16,
          num_classes=10,
          num_stages=2,
          embed_dims=[2, 4],
          num_heads=[2, 4],
          mlp_ratios=[2, 2],
          depths=[2, 3],
          sr_ratios=[2, 1],
      ).to(self.device)
      outs1 = model1(dummy_img)
      self.assertTrue(outs1.shape == (1, 10))
  
      model2 = create_model(
          img_size=16,
          num_classes=0,
          num_stages=2,
          embed_dims=[2, 4],
          num_heads=[2, 4],
          mlp_ratios=[2, 2],
          depths=[2, 3],
          sr_ratios=[2, 1],
      ).to(self.device)
      outs2 = model2(dummy_img)
      self.assertTrue(outs2.shape == (1, 4))
  
imports: import unittest
  import torch
  from towhee.models.shunted_transformer import create_model
label_desc: Use image node_id 1 for calling the TestModel obj's underlying member method code with example usage: obj.test_base() without return types