def test_inference_fp16(self):
        r"""
        A small test to make sure that inference work in half precision without any problem.
        """
        model = DeiTModel.from_pretrained(
            "facebook/deit-base-distilled-patch16-224", torch_dtype=torch.float16, device_map="auto"
        )
        image_processor = self.default_image_processor

        image = prepare_img()
        inputs = image_processor(images=image, return_tensors="pt")
        pixel_values = inputs.pixel_values.to(torch_device)

        # forward pass to make sure inference works in fp16
        with torch.no_grad():
            _ = model(pixel_values)

        A small test to make sure that inference work in half precision without any problem.
        
dpkd/transformers/tests/models/deit/test_modeling_deit.py