def METHOD_NAME(self):
    codec_cfg = dict(
        type='AssociativeEmbedding',
        input_size=(256, 256),
        heatmap_size=(64, 64),
        use_udp=False,
        decode_keypoint_order=[
            0, 1, 2, 3, 4, 5, 6, 11, 12, 7, 8, 9, 10, 13, 14, 15, 16
        ])
    # get test data
    codec = KEYPOINT_CODECS.build(codec_cfg)
    batch_data_samples = get_packed_inputs(
        1,
        input_size=(256, 256),
        heatmap_size=(64, 64),
        img_shape=(256, 256))['data_samples']
    keypoints = batch_data_samples[0].gt_instances['keypoints']
    keypoints_visible = batch_data_samples[0].gt_instances[
        'keypoints_visible']
    encoded = codec.encode(keypoints, keypoints_visible)
    heatmaps = encoded['heatmaps']
    keypoint_indices = encoded['keypoint_indices']
    keypoint_weights = encoded['keypoint_weights']
    heatmap_mask = np.ones((1, ) + heatmaps.shape[1:], dtype=np.float32)
    batch_data_samples[0].gt_fields = PixelData(
        heatmaps=heatmaps, heatmap_mask=heatmap_mask).to_tensor()
    batch_data_samples[0].gt_instance_labels = InstanceData(
        keypoint_indices=keypoint_indices,
        keypoint_weights=keypoint_weights).to_tensor()
    feats = [torch.rand(1, 32, 64, 64)]
    head = AssociativeEmbeddingHead(
        in_channels=32,
        num_keypoints=17,
        tag_dim=1,
        tag_per_keypoint=True,
        deconv_out_channels=None)
    losses = head.loss(feats, batch_data_samples)
    for name in ['loss_kpt', 'loss_pull', 'loss_push']:
        self.assertIn(name, losses)
        self.assertIsInstance(losses[name], torch.Tensor)