def METHOD_NAME(self):
    self.model_config["training_head_type"] = "classification"
    finetune_model = build_model(self.model_config)
    finetune_model.model.eval()
    num_bbox_per_image = 10
    input_ids = torch.randint(low=0, high=BERT_VOCAB_SIZE, size=(1, 128)).long()
    attention_mask = torch.ones((1, 128)).long()
    token_type_ids = torch.zeros(1, 128).long()
    visual_embeddings = torch.rand(
        (1, num_bbox_per_image, self.vision_feature_size)
    ).float()
    image_attention_mask = torch.zeros((1, num_bbox_per_image)).long()
    visual_locations = torch.rand((1, num_bbox_per_image, 5)).float()
    finetune_model.eval()
    with torch.no_grad():
        model_output = finetune_model.model(
            input_ids=input_ids,
            image_feature=visual_embeddings,
            image_location=visual_locations,
            token_type_ids=token_type_ids,
            attention_mask=attention_mask,
            image_attention_mask=image_attention_mask,
        )
    script_model = torch.jit.script(finetune_model.model)
    with torch.no_grad():
        script_output = script_model(
            input_ids=input_ids,
            image_feature=visual_embeddings,
            image_location=visual_locations,
            token_type_ids=token_type_ids,
            attention_mask=attention_mask,
            image_attention_mask=image_attention_mask,
        )
    self.assertTrue(torch.equal(model_output["scores"], script_output["scores"]))