def METHOD_NAME(self):
    extra = dict(
        stage1=dict(
            num_modules=1,
            num_branches=1,
            block='BOTTLENECK',
            num_blocks=(4, ),
            num_channels=(64, )),
        stage2=dict(
            num_modules=1,
            num_branches=2,
            block='BASIC',
            num_blocks=(4, 4),
            num_channels=(32, 64)),
        stage3=dict(
            num_modules=4,
            num_branches=3,
            block='BASIC',
            num_blocks=(4, 4, 4),
            num_channels=(32, 64, 128)),
        stage4=dict(
            num_modules=3,
            num_branches=4,
            block='BASIC',
            num_blocks=(4, 4, 4, 4),
            num_channels=(32, 64, 128, 256)))
    model = HRNet(extra, in_channels=3)
    imgs = torch.randn(2, 3, 224, 224)
    feat = model(imgs)
    self.assertIsInstance(feat, tuple)
    self.assertEqual(feat[-1].shape, torch.Size([2, 32, 56, 56]))
    # Test HRNet zero initialization of residual
    model = HRNet(extra, in_channels=3, zero_init_residual=True)
    model.init_weights()
    for m in model.modules():
        if isinstance(m, Bottleneck):
            self.assertTrue(self.all_zeros(m.norm3))
    model.train()
    imgs = torch.randn(2, 3, 224, 224)
    feat = model(imgs)
    self.assertIsInstance(feat, tuple)
    self.assertEqual(feat[-1].shape, torch.Size([2, 32, 56, 56]))
    # Test HRNet with the first three stages frozen
    frozen_stages = 3
    model = HRNet(extra, in_channels=3, frozen_stages=frozen_stages)
    model.init_weights()
    model.train()
    if frozen_stages >= 0:
        self.assertFalse(model.norm1.training)
        self.assertFalse(model.norm2.training)
        for layer in [model.conv1, model.norm1, model.conv2, model.norm2]:
            for param in layer.parameters():
                self.assertFalse(param.requires_grad)
    for i in range(1, frozen_stages + 1):
        if i == 1:
            layer = getattr(model, 'layer1')
        else:
            layer = getattr(model, f'stage{i}')
        for mod in layer.modules():
            if isinstance(mod, _BatchNorm):
                self.assertFalse(mod.training)
        for param in layer.parameters():
            self.assertFalse(param.requires_grad)
        if i < 4:
            layer = getattr(model, f'transition{i}')
        for mod in layer.modules():
            if isinstance(mod, _BatchNorm):
                self.assertFalse(mod.training)
        for param in layer.parameters():
            self.assertFalse(param.requires_grad)