Note that this model does not work directly with HF, a modification that does mean pooling before the layernorm and classification head is needed.


```python
from transformers import (
    ViTForImageClassification,
    pipeline,
    AutoImageProcessor,
    ViTConfig,
    ViTModel,
)

from transformers.modeling_outputs import (
    ImageClassifierOutput,
    BaseModelOutputWithPooling,
)

from PIL import Image
import torch
from torch import nn
from typing import Optional, Union, Tuple


class CustomViTModel(ViTModel):
    def forward(
        self,
        pixel_values: Optional[torch.Tensor] = None,
        bool_masked_pos: Optional[torch.BoolTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        interpolate_pos_encoding: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r"""
        bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*):
            Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
        """
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        # TODO: maybe have a cleaner way to cast the input (from `ImageProcessor` side?)
        expected_dtype = self.embeddings.patch_embeddings.projection.weight.dtype
        if pixel_values.dtype != expected_dtype:
            pixel_values = pixel_values.to(expected_dtype)

        embedding_output = self.embeddings(
            pixel_values,
            bool_masked_pos=bool_masked_pos,
            interpolate_pos_encoding=interpolate_pos_encoding,
        )

        encoder_outputs = self.encoder(
            embedding_output,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        sequence_output = sequence_output[:, 1:, :].mean(dim=1)

        sequence_output = self.layernorm(sequence_output)
        pooled_output = (
            self.pooler(sequence_output) if self.pooler is not None else None
        )

        if not return_dict:
            head_outputs = (
                (sequence_output, pooled_output)
                if pooled_output is not None
                else (sequence_output,)
            )
            return head_outputs + encoder_outputs[1:]

        return BaseModelOutputWithPooling(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )


class CustomViTForImageClassification(ViTForImageClassification):
    def __init__(self, config: ViTConfig) -> None:
        super().__init__(config)

        self.num_labels = config.num_labels
        self.vit = CustomViTModel(config, add_pooling_layer=False)

        # Classifier head
        self.classifier = (
            nn.Linear(config.hidden_size, config.num_labels)
            if config.num_labels > 0
            else nn.Identity()
        )

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        pixel_values: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        interpolate_pos_encoding: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[tuple, ImageClassifierOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        """
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        outputs = self.vit(
            pixel_values,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            interpolate_pos_encoding=interpolate_pos_encoding,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]

        logits = self.classifier(sequence_output)

        loss = None

        return ImageClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


model = CustomViTForImageClassification.from_pretrained("vesteinn/vit-mae-cub")
image_processor = AutoImageProcessor.from_pretrained("vesteinn/vit-mae-cub")

classifier = pipeline(
    "image-classification", model=model, image_processor=image_processor
)
```
black footed albatross laysan albatross sooty albatross groove billed ani crested auklet least auklet parakeet auklet rhinoceros auklet brewer blackbird red winged blackbird rusty blackbird yellow headed blackbird bobolink indigo bunting lazuli bunting painted bunting cardinal spotted catbird gray catbird yellow breasted chat eastern towhee chuck will widow brandt cormorant red faced cormorant pelagic cormorant bronzed cowbird shiny cowbird brown creeper american crow fish crow black billed cuckoo mangrove cuckoo yellow billed cuckoo gray crowned rosy finch purple finch northern flicker acadian flycatcher great crested flycatcher least flycatcher olive sided flycatcher scissor tailed flycatcher vermilion flycatcher yellow bellied flycatcher frigatebird northern fulmar gadwall american goldfinch european goldfinch boat tailed grackle eared grebe horned grebe pied billed grebe western grebe blue grosbeak evening grosbeak pine grosbeak rose breasted grosbeak pigeon guillemot california gull glaucous winged gull heermann gull herring gull ivory gull ring billed gull slaty backed gull western gull anna hummingbird ruby throated hummingbird rufous hummingbird green violetear long tailed jaeger pomarine jaeger blue jay florida jay green jay dark eyed junco tropical kingbird gray kingbird belted kingfisher green kingfisher pied kingfisher ringed kingfisher white breasted kingfisher red legged kittiwake horned lark pacific loon mallard western meadowlark hooded merganser red breasted merganser mockingbird nighthawk clark nutcracker white breasted nuthatch baltimore oriole hooded oriole orchard oriole scott oriole ovenbird brown pelican white pelican western wood pewee sayornis american pipit whip poor will horned puffin common raven white necked raven american redstart geococcyx loggerhead shrike great grey shrike baird sparrow black throated sparrow brewer sparrow chipping sparrow clay colored sparrow house sparrow field sparrow fox sparrow grasshopper sparrow harris sparrow henslow sparrow le conte sparrow lincoln sparrow nelson sharp tailed sparrow savannah sparrow seaside sparrow song sparrow tree sparrow vesper sparrow white crowned sparrow white throated sparrow cape glossy starling bank swallow barn swallow cliff swallow tree swallow scarlet tanager summer tanager artic tern black tern caspian tern common tern elegant tern forsters tern least tern green tailed towhee brown thrasher sage thrasher black capped vireo blue headed vireo philadelphia vireo red eyed vireo warbling vireo white eyed vireo yellow throated vireo bay breasted warbler black and white warbler black throated blue warbler blue winged warbler canada warbler cape may warbler cerulean warbler chestnut sided warbler golden winged warbler hooded warbler kentucky warbler magnolia warbler mourning warbler myrtle warbler nashville warbler orange crowned warbler palm warbler pine warbler prairie warbler prothonotary warbler swainson warbler tennessee warbler wilson warbler worm eating warbler yellow warbler northern waterthrush louisiana waterthrush bohemian waxwing cedar waxwing american three toed woodpecker pileated woodpecker red bellied woodpecker red cockaded woodpecker red headed woodpecker downy woodpecker bewick wren cactus wren carolina wren house wren marsh wren rock wren winter wren common yellowthroat label_200 label_201 label_202 label_203 label_204 label_205 label_206 label_207 label_208 label_209 label_210 label_211 label_212 label_213 label_214 label_215 label_216 label_217 label_218 label_219 label_220 label_221 label_222 label_223 label_224 label_225 label_226 label_227 label_228 label_229 label_230 label_231 label_232 label_233 label_234 label_235 label_236 label_237 label_238 label_239 label_240 label_241 label_242 label_243 label_244 label_245 label_246 label_247 label_248 label_249 label_250 label_251 label_252 label_253 label_254 label_255 label_256 label_257 label_258 label_259 label_260 label_261 label_262 label_263 label_264 label_265 label_266 label_267 label_268 label_269 label_270 label_271 label_272 label_273 label_274 label_275 label_276 label_277 label_278 label_279 label_280 label_281 label_282 label_283 label_284 label_285 label_286 label_287 label_288 label_289 label_290 label_291 label_292 label_293 label_294 label_295 label_296 label_297 label_298 label_299 label_300 label_301 label_302 label_303 label_304 label_305 label_306 label_307 label_308 label_309 label_310 label_311 label_312 label_313 label_314 label_315 label_316 label_317 label_318 label_319 label_320 label_321 label_322 label_323 label_324 label_325 label_326 label_327 label_328 label_329 label_330 label_331 label_332 label_333 label_334 label_335 label_336 label_337 label_338 label_339 label_340 label_341 label_342 label_343 label_344 label_345 label_346 label_347 label_348 label_349 label_350 label_351 label_352 label_353 label_354 label_355 label_356 label_357 label_358 label_359 label_360 label_361 label_362 label_363 label_364 label_365 label_366 label_367 label_368 label_369 label_370 label_371 label_372 label_373 label_374 label_375 label_376 label_377 label_378 label_379 label_380 label_381 label_382 label_383 label_384 label_385 label_386 label_387 label_388 label_389 label_390 label_391 label_392 label_393 label_394 label_395 label_396 label_397 label_398 label_399 label_400 label_401 label_402 label_403 label_404 label_405 label_406 label_407 label_408 label_409 label_410 label_411 label_412 label_413 label_414 label_415 label_416 label_417 label_418 label_419 label_420 label_421 label_422 label_423 label_424 label_425 label_426 label_427 label_428 label_429 label_430 label_431 label_432 label_433 label_434 label_435 label_436 label_437 label_438 label_439 label_440 label_441 label_442 label_443 label_444 label_445 label_446 label_447 label_448 label_449 label_450 label_451 label_452 label_453 label_454 label_455 label_456 label_457 label_458 label_459 label_460 label_461 label_462 label_463 label_464 label_465 label_466 label_467 label_468 label_469 label_470 label_471 label_472 label_473 label_474 label_475 label_476 label_477 label_478 label_479 label_480 label_481 label_482 label_483 label_484 label_485 label_486 label_487 label_488 label_489 label_490 label_491 label_492 label_493 label_494 label_495 label_496 label_497 label_498 label_499 label_500 label_501 label_502 label_503 label_504 label_505 label_506 label_507 label_508 label_509 label_510 label_511 label_512 label_513 label_514 label_515 label_516 label_517 label_518 label_519 label_520 label_521 label_522 label_523 label_524 label_525 label_526 label_527 label_528 label_529 label_530 label_531 label_532 label_533 label_534 label_535 label_536 label_537 label_538 label_539 label_540 label_541 label_542 label_543 label_544 label_545 label_546 label_547 label_548 label_549 label_550 label_551 label_552 label_553 label_554 label_555 label_556 label_557 label_558 label_559 label_560 label_561 label_562 label_563 label_564 label_565 label_566 label_567 label_568 label_569 label_570 label_571 label_572 label_573 label_574 label_575 label_576 label_577 label_578 label_579 label_580 label_581 label_582 label_583 label_584 label_585 label_586 label_587 label_588 label_589 label_590 label_591 label_592 label_593 label_594 label_595 label_596 label_597 label_598 label_599 label_600 label_601 label_602 label_603 label_604 label_605 label_606 label_607 label_608 label_609 label_610 label_611 label_612 label_613 label_614 label_615 label_616 label_617 label_618 label_619 label_620 label_621 label_622 label_623 label_624 label_625 label_626 label_627 label_628 label_629 label_630 label_631 label_632 label_633 label_634 label_635 label_636 label_637 label_638 label_639 label_640 label_641 label_642 label_643 label_644 label_645 label_646 label_647 label_648 label_649 label_650 label_651 label_652 label_653 label_654 label_655 label_656 label_657 label_658 label_659 label_660 label_661 label_662 label_663 label_664 label_665 label_666 label_667 label_668 label_669 label_670 label_671 label_672 label_673 label_674 label_675 label_676 label_677 label_678 label_679 label_680 label_681 label_682 label_683 label_684 label_685 label_686 label_687 label_688 label_689 label_690 label_691 label_692 label_693 label_694 label_695 label_696 label_697 label_698 label_699 label_700 label_701 label_702 label_703 label_704 label_705 label_706 label_707 label_708 label_709 label_710 label_711 label_712 label_713 label_714 label_715 label_716 label_717 label_718 label_719 label_720 label_721 label_722 label_723 label_724 label_725 label_726 label_727 label_728 label_729 label_730 label_731 label_732 label_733 label_734 label_735 label_736 label_737 label_738 label_739 label_740 label_741 label_742 label_743 label_744 label_745 label_746 label_747 label_748 label_749 label_750 label_751 label_752 label_753 label_754 label_755 label_756 label_757 label_758 label_759 label_760 label_761 label_762 label_763 label_764 label_765 label_766 label_767 label_768 label_769 label_770 label_771 label_772 label_773 label_774 label_775 label_776 label_777 label_778 label_779 label_780 label_781 label_782 label_783 label_784 label_785 label_786 label_787 label_788 label_789 label_790 label_791 label_792 label_793 label_794 label_795 label_796 label_797 label_798 label_799 label_800 label_801 label_802 label_803 label_804 label_805 label_806 label_807 label_808 label_809 label_810 label_811 label_812 label_813 label_814 label_815 label_816 label_817 label_818 label_819 label_820 label_821 label_822 label_823 label_824 label_825 label_826 label_827 label_828 label_829 label_830 label_831 label_832 label_833 label_834 label_835 label_836 label_837 label_838 label_839 label_840 label_841 label_842 label_843 label_844 label_845 label_846 label_847 label_848 label_849 label_850 label_851 label_852 label_853 label_854 label_855 label_856 label_857 label_858 label_859 label_860 label_861 label_862 label_863 label_864 label_865 label_866 label_867 label_868 label_869 label_870 label_871 label_872 label_873 label_874 label_875 label_876 label_877 label_878 label_879 label_880 label_881 label_882 label_883 label_884 label_885 label_886 label_887 label_888 label_889 label_890 label_891 label_892 label_893 label_894 label_895 label_896 label_897 label_898 label_899 label_900 label_901 label_902 label_903 label_904 label_905 label_906 label_907 label_908 label_909 label_910 label_911 label_912 label_913 label_914 label_915 label_916 label_917 label_918 label_919 label_920 label_921 label_922 label_923 label_924 label_925 label_926 label_927 label_928 label_929 label_930 label_931 label_932 label_933 label_934 label_935 label_936 label_937 label_938 label_939 label_940 label_941 label_942 label_943 label_944 label_945 label_946 label_947 label_948 label_949 label_950 label_951 label_952 label_953 label_954 label_955 label_956 label_957 label_958 label_959 label_960 label_961 label_962 label_963 label_964 label_965 label_966 label_967 label_968 label_969 label_970 label_971 label_972 label_973 label_974 label_975 label_976 label_977 label_978 label_979 label_980 label_981 label_982 label_983 label_984 label_985 label_986 label_987 label_988 label_989 label_990 label_991 label_992 label_993 label_994 label_995 label_996 label_997 label_998 label_999