def METHOD_NAME(self):
    norm_cfg = dict(type='BN')
    # only have 3 stages
    extra = dict(
        drop_path_rate=0.2,
        stage1=dict(
            num_modules=1,
            num_branches=1,
            block='BOTTLENECK',
            num_blocks=(2, ),
            num_channels=(64, )),
        stage2=dict(
            num_modules=1,
            num_branches=2,
            block='HRFORMERBLOCK',
            window_sizes=(7, 7),
            num_heads=(1, 2),
            mlp_ratios=(4, 4),
            num_blocks=(2, 2),
            num_channels=(32, 64)),
        stage3=dict(
            num_modules=4,
            num_branches=3,
            block='HRFORMERBLOCK',
            window_sizes=(7, 7, 7),
            num_heads=(1, 2, 4),
            mlp_ratios=(4, 4, 4),
            num_blocks=(2, 2, 2),
            num_channels=(32, 64, 128)),
        stage4=dict(
            num_modules=3,
            num_branches=4,
            block='HRFORMERBLOCK',
            window_sizes=(7, 7, 7, 7),
            num_heads=(1, 2, 4, 8),
            mlp_ratios=(4, 4, 4, 4),
            num_blocks=(2, 2, 2, 2),
            num_channels=(32, 64, 128, 256),
            multiscale_output=True))
    with self.assertRaises(ValueError):
        # len(num_blocks) should equal num_branches
        extra['stage4']['num_branches'] = 3
        HRFormer(extra=extra)
    extra['stage4']['num_branches'] = 4
    # Test HRFormer-S
    model = HRFormer(extra=extra, norm_cfg=norm_cfg)
    model.init_weights()
    model.train()
    imgs = torch.randn(1, 3, 64, 64)
    feats = model(imgs)
    self.assertEqual(len(feats), 4)
    self.assertEqual(feats[0].shape, torch.Size([1, 32, 16, 16]))
    self.assertEqual(feats[3].shape, torch.Size([1, 256, 2, 2]))
    # Test single scale output and model
    # without relative position bias
    extra['stage4']['multiscale_output'] = False
    extra['with_rpe'] = False
    model = HRFormer(extra=extra, norm_cfg=norm_cfg)
    model.init_weights()
    model.train()
    imgs = torch.randn(1, 3, 64, 64)
    feats = model(imgs)
    self.assertIsInstance(feats, tuple)
    self.assertEqual(len(feats), 1)
    self.assertEqual(feats[-1].shape, torch.Size([1, 32, 16, 16]))