class_docstr: Test utils for Shunted Transformer
code: def test_attention(self):
      with self.assertRaises(AssertionError) as e:
          Attention(dim=4, num_heads=3, qk_scale=1).to(self.device)
      self.assertEqual(
          str(e.exception), "dim 4 should be divisible by num_heads 3."
      )
  
      input1 = torch.rand(1, 4, 4).to(self.device)
      model1 = Attention(dim=4, num_heads=4, qk_scale=1, sr_ratio=2).to(
          self.device
      )
      out1 = model1(input1, 2, 2)
      self.assertTrue(out1.shape == (1, 4, 4))
  
      input2 = torch.rand(1, 4, 6).to(self.device)
      model2 = Attention(dim=6, num_heads=3, qkv_bias=True).to(
          self.device
      )
      out2 = model2(input2, 2, 2)
      self.assertTrue(out2.shape == (1, 4, 6))
  
file_path: towhee/tests/unittests/models/shunted_transformer/test_utils.py
imports: import unittest
  import torch
  from towhee.models.shunted_transformer import OverlapPatchEmbed, Mlp, Attention, Block, HeadPatchEmbed
label_desc: Use image node_id 4 for calling the TestUtils obj's underlying member method code with example usage: obj.test_attention() without return types