






import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DeiTConfig, DeiTForImageClassificationWithTeacher, DeiTImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase : Tuple						   =       logging.get_logger(__name__)
def       snake_case_		(     lowerCAmelCase_       :						Dict       ,       lowerCAmelCase_       :						Union[str, Any]=False      ):
							__lowercase      :							Optional[Any]				=     []
							for i in range(config.num_hidden_layers      ):
														# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
														rename_keys.append((F"blocks.{i}.norm1.weight", F"deit.encoder.layer.{i}.layernorm_before.weight")      )
														rename_keys.append((F"blocks.{i}.norm1.bias", F"deit.encoder.layer.{i}.layernorm_before.bias")      )
														rename_keys.append((F"blocks.{i}.attn.proj.weight", F"deit.encoder.layer.{i}.attention.output.dense.weight")      )
														rename_keys.append((F"blocks.{i}.attn.proj.bias", F"deit.encoder.layer.{i}.attention.output.dense.bias")      )
														rename_keys.append((F"blocks.{i}.norm2.weight", F"deit.encoder.layer.{i}.layernorm_after.weight")      )
														rename_keys.append((F"blocks.{i}.norm2.bias", F"deit.encoder.layer.{i}.layernorm_after.bias")      )
														rename_keys.append((F"blocks.{i}.mlp.fc1.weight", F"deit.encoder.layer.{i}.intermediate.dense.weight")      )
														rename_keys.append((F"blocks.{i}.mlp.fc1.bias", F"deit.encoder.layer.{i}.intermediate.dense.bias")      )
														rename_keys.append((F"blocks.{i}.mlp.fc2.weight", F"deit.encoder.layer.{i}.output.dense.weight")      )
														rename_keys.append((F"blocks.{i}.mlp.fc2.bias", F"deit.encoder.layer.{i}.output.dense.bias")      )

							# projection layer + position embeddings
							rename_keys.extend(
							    [
							        ("""cls_token""", """deit.embeddings.cls_token"""),
							        ("""dist_token""", """deit.embeddings.distillation_token"""),
							        ("""patch_embed.proj.weight""", """deit.embeddings.patch_embeddings.projection.weight"""),
							        ("""patch_embed.proj.bias""", """deit.embeddings.patch_embeddings.projection.bias"""),
							        ("""pos_embed""", """deit.embeddings.position_embeddings"""),
							    ]      )

							if base_model:
														# layernorm + pooler
														rename_keys.extend(
														    [
														        ("""norm.weight""", """layernorm.weight"""),
														        ("""norm.bias""", """layernorm.bias"""),
														        ("""pre_logits.fc.weight""", """pooler.dense.weight"""),
														        ("""pre_logits.fc.bias""", """pooler.dense.bias"""),
														    ]      )

														# if just the base model, we should remove "deit" from all keys that start with "deit"
														__lowercase      :							Optional[int]				=     [(pair[0], pair[1][4:]) if pair[1].startswith("""deit"""      ) else pair for pair in rename_keys]
							else:
														# layernorm + classification heads
														rename_keys.extend(
														    [
														        ("""norm.weight""", """deit.layernorm.weight"""),
														        ("""norm.bias""", """deit.layernorm.bias"""),
														        ("""head.weight""", """cls_classifier.weight"""),
														        ("""head.bias""", """cls_classifier.bias"""),
														        ("""head_dist.weight""", """distillation_classifier.weight"""),
														        ("""head_dist.bias""", """distillation_classifier.bias"""),
														    ]      )

							return rename_keys
def       snake_case_		(     lowerCAmelCase_       :						List[str]       ,       lowerCAmelCase_       :						Optional[Any]       ,       lowerCAmelCase_       :						str=False      ):
							for i in range(config.num_hidden_layers      ):
														if base_model:
																					__lowercase      :							List[str]				=     """"""
														else:
																					__lowercase      :							Optional[Any]				=     """deit."""
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														__lowercase      :							Dict				=     state_dict.pop(F"blocks.{i}.attn.qkv.weight"      )
														__lowercase      :							str				=     state_dict.pop(F"blocks.{i}.attn.qkv.bias"      )
														# next, add query, keys and values (in that order) to the state dict
														__lowercase      :							Optional[int]				=     in_proj_weight[
														    : config.hidden_size, :
														]
														__lowercase      :							Optional[int]				=     in_proj_bias[: config.hidden_size]
														__lowercase      :							str				=     in_proj_weight[
														    config.hidden_size : config.hidden_size * 2, :
														]
														__lowercase      :							Tuple				=     in_proj_bias[
														    config.hidden_size : config.hidden_size * 2
														]
														__lowercase      :							Dict				=     in_proj_weight[
														    -config.hidden_size :, :
														]
														__lowercase      :							Tuple				=     in_proj_bias[-config.hidden_size :]
def       snake_case_		(     lowerCAmelCase_       :						Dict       ,       lowerCAmelCase_       :						Dict       ,       lowerCAmelCase_       :						int      ):
							__lowercase      :							Union[str, Any]				=     dct.pop(lowerCAmelCase_      )
							__lowercase      :							Union[str, Any]				=     val
def       snake_case_		(     ):
							__lowercase      :							Optional[Any]				=     """http://images.cocodataset.org/val2017/000000039769.jpg"""
							__lowercase      :							List[Any]				=     Image.open(requests.get(lowerCAmelCase_       ,       stream=lowerCAmelCase_      ).raw      )
							return im
@torch.no_grad()
def       snake_case_		(     lowerCAmelCase_       :						List[str]       ,       lowerCAmelCase_       :						str      ):
							__lowercase      :							Tuple				=     DeiTConfig()
							# all deit models have fine-tuned heads
							__lowercase      :							Dict				=     False
							# dataset (fine-tuned on ImageNet 2012), patch_size and image_size
							__lowercase      :							Any				=     1000
							__lowercase      :							str				=     """huggingface/label-files"""
							__lowercase      :							Union[str, Any]				=     """imagenet-1k-id2label.json"""
							__lowercase      :							Any				=     json.load(open(hf_hub_download(lowerCAmelCase_       ,       lowerCAmelCase_       ,       repo_type="""dataset"""      )       ,       """r"""      )      )
							__lowercase      :							Dict				=     {int(lowerCAmelCase_      ): v for k, v in idalabel.items()}
							__lowercase      :							int				=     idalabel
							__lowercase      :							Optional[Any]				=     {v: k for k, v in idalabel.items()}
							__lowercase      :							Union[str, Any]				=     int(deit_name[-6:-4]      )
							__lowercase      :							Any				=     int(deit_name[-3:]      )
							# size of the architecture
							if deit_name[9:].startswith("""tiny"""      ):
														__lowercase      :							Tuple				=     192
														__lowercase      :							Optional[int]				=     768
														__lowercase      :							str				=     12
														__lowercase      :							Tuple				=     3
							elif deit_name[9:].startswith("""small"""      ):
														__lowercase      :							List[str]				=     384
														__lowercase      :							Tuple				=     1536
														__lowercase      :							List[str]				=     12
														__lowercase      :							Optional[Any]				=     6
							if deit_name[9:].startswith("""base"""      ):
														pass
							elif deit_name[4:].startswith("""large"""      ):
														__lowercase      :							List[Any]				=     1024
														__lowercase      :							Optional[Any]				=     4096
														__lowercase      :							Optional[Any]				=     24
														__lowercase      :							Optional[Any]				=     16

							# load original model from timm
							__lowercase      :							int				=     timm.create_model(lowerCAmelCase_       ,       pretrained=lowerCAmelCase_      )
							timm_model.eval()

							# load state_dict of original model, remove and rename some keys
							__lowercase      :							Dict				=     timm_model.state_dict()
							__lowercase      :							List[str]				=     create_rename_keys(lowerCAmelCase_       ,       lowerCAmelCase_      )
							for src, dest in rename_keys:
														rename_key(lowerCAmelCase_       ,       lowerCAmelCase_       ,       lowerCAmelCase_      )
							read_in_q_k_v(lowerCAmelCase_       ,       lowerCAmelCase_       ,       lowerCAmelCase_      )

							# load HuggingFace model
							__lowercase      :							int				=     DeiTForImageClassificationWithTeacher(lowerCAmelCase_      ).eval()
							model.load_state_dict(lowerCAmelCase_      )

							# Check outputs on an image, prepared by DeiTImageProcessor
							__lowercase      :							Tuple				=     int(
							    (256 / 224) * config.image_size      )  # to maintain same ratio w.r.t. 224 images, see https://github.com/facebookresearch/deit/blob/ab5715372db8c6cad5740714b2216d55aeae052e/datasets.py#L103
							__lowercase      :							int				=     DeiTImageProcessor(size=lowerCAmelCase_       ,       crop_size=config.image_size      )
							__lowercase      :							Union[str, Any]				=     image_processor(images=prepare_img()       ,       return_tensors="""pt"""      )
							__lowercase      :							List[Any]				=     encoding["""pixel_values"""]
							__lowercase      :							Any				=     model(lowerCAmelCase_      )

							__lowercase      :							Optional[int]				=     timm_model(lowerCAmelCase_      )
							assert timm_logits.shape == outputs.logits.shape
							assert torch.allclose(lowerCAmelCase_       ,       outputs.logits       ,       atol=1e-3      )

							Path(lowerCAmelCase_      ).mkdir(exist_ok=lowerCAmelCase_      )
							print(F"Saving model {deit_name} to {pytorch_dump_folder_path}"      )
							model.save_pretrained(lowerCAmelCase_      )
							print(F"Saving image processor to {pytorch_dump_folder_path}"      )
							image_processor.save_pretrained(lowerCAmelCase_      )


if __name__ == "__main__":
	lowerCamelCase : Dict						   =       argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '''--deit_name''',
	    default='''vit_deit_base_distilled_patch16_224''',
	    type=str,
	    help='''Name of the DeiT timm model you\'d like to convert.''',
	)
	parser.add_argument(
	    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
	)

	lowerCamelCase : List[Any]						   =       parser.parse_args()
	convert_deit_checkpoint(args.deit_name, args.pytorch_dump_folder_path)







import inspect
import unittest

from transformers import MobileNetVaConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import MobileNetVaForImageClassification, MobileNetVaForSemanticSegmentation, MobileNetVaModel
	from transformers.models.mobilenet_va.modeling_mobilenet_va import MOBILENET_V2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import MobileNetVaImageProcessor

class 		lowerCAmelCase     ( __a			):

					'''simple docstring'''

					def       lowerCAmelCase				(  self				:	List[str]							)			->  Optional[int]:


												"""simple docstring"""

												__lowercase      :							str				=     self.config_class(**self.inputs_dict							)
												self.parent.assertTrue(hasattr(__a		,     """tf_padding"""							)							)
												self.parent.assertTrue(hasattr(__a		,     """depth_multiplier"""							)							)







class 		lowerCAmelCase     :

					'''simple docstring'''

					def __init__(  self				:	Union[str, Any]		,     __a				:	Tuple		,     __a				:	str=13		,     __a				:	Dict=3		,     __a				:	List[Any]=32		,     __a				:	Any=0.25		,     __a				:	Any=8		,     __a				:	Optional[int]=8		,     __a				:	Optional[int]=6		,     __a				:	Dict=32		,     __a				:	Tuple=True		,     __a				:	List[Any]=True		,     __a				:	Optional[int]=True		,     __a				:	Tuple="relu6"		,     __a				:	Optional[Any]=1280		,     __a				:	str=0.1		,     __a				:	str=0.02		,     __a				:	Optional[Any]=True		,     __a				:	Tuple=True		,     __a				:	Dict=10		,     __a				:	Optional[Any]=None		,     )			->  Any:


												"""simple docstring"""

												__lowercase      :							List[str]				=     parent
												__lowercase      :							Tuple				=     batch_size
												__lowercase      :							Dict				=     num_channels
												__lowercase      :							Optional[int]				=     image_size
												__lowercase      :							int				=     depth_multiplier
												__lowercase      :							str				=     depth_divisible_by
												__lowercase      :							int				=     min_depth
												__lowercase      :							Tuple				=     expand_ratio
												__lowercase      :							Optional[int]				=     tf_padding
												__lowercase      :							Dict				=     output_stride
												__lowercase      :							Dict				=     first_layer_is_expansion
												__lowercase      :							Optional[Any]				=     finegrained_output
												__lowercase      :							str				=     hidden_act
												__lowercase      :							Union[str, Any]				=     last_hidden_size if finegrained_output else int(last_hidden_size * depth_multiplier							)
												__lowercase      :							Optional[int]				=     classifier_dropout_prob
												__lowercase      :							int				=     use_labels
												__lowercase      :							Optional[int]				=     is_training
												__lowercase      :							Dict				=     num_labels
												__lowercase      :							Tuple				=     initializer_range
												__lowercase      :							Optional[Any]				=     scope



					def       lowerCAmelCase				(  self				:	Any							)			->  Optional[Any]:


												"""simple docstring"""

												__lowercase      :							Tuple				=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

												__lowercase      :							List[Any]				=     None
												__lowercase      :							Optional[Any]				=     None
												if self.use_labels:
																			__lowercase      :							List[Any]				=     ids_tensor([self.batch_size]		,     self.num_labels							)
																			__lowercase      :							Optional[int]				=     ids_tensor([self.batch_size, self.image_size, self.image_size]		,     self.num_labels							)

												__lowercase      :							List[Any]				=     self.get_config()

												return config, pixel_values, labels, pixel_labels



					def       lowerCAmelCase				(  self				:	str							)			->  Union[str, Any]:


												"""simple docstring"""

												return MobileNetVaConfig(
												    num_channels=self.num_channels		,     image_size=self.image_size		,     depth_multiplier=self.depth_multiplier		,     depth_divisible_by=self.depth_divisible_by		,     min_depth=self.min_depth		,     expand_ratio=self.expand_ratio		,     output_stride=self.output_stride		,     first_layer_is_expansion=self.first_layer_is_expansion		,     finegrained_output=self.finegrained_output		,     hidden_act=self.hidden_act		,     tf_padding=self.tf_padding		,     classifier_dropout_prob=self.classifier_dropout_prob		,     initializer_range=self.initializer_range		,     )



					def       lowerCAmelCase				(  self				:	Tuple		,     __a				:	Dict		,     __a				:	Tuple		,     __a				:	Optional[int]		,     __a				:	Union[str, Any]							)			->  List[Any]:


												"""simple docstring"""

												__lowercase      :							Optional[int]				=     MobileNetVaModel(config=__a							)
												model.to(__a							)
												model.eval()
												__lowercase      :							Tuple				=     model(__a							)
												self.parent.assertEqual(
												    result.last_hidden_state.shape		,     (
												        self.batch_size,
												        self.last_hidden_size,
												        self.image_size // self.output_stride,
												        self.image_size // self.output_stride,
												    )		,     )
												self.parent.assertEqual(
												    result.pooler_output.shape		,     (self.batch_size, self.last_hidden_size)		,     )



					def       lowerCAmelCase				(  self				:	List[str]		,     __a				:	Optional[int]		,     __a				:	List[str]		,     __a				:	str		,     __a				:	Optional[int]							)			->  Tuple:


												"""simple docstring"""

												__lowercase      :							List[Any]				=     self.num_labels
												__lowercase      :							Dict				=     MobileNetVaForImageClassification(__a							)
												model.to(__a							)
												model.eval()
												__lowercase      :							Dict				=     model(__a		,     labels=__a							)
												self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.num_labels)							)



					def       lowerCAmelCase				(  self				:	int		,     __a				:	List[str]		,     __a				:	Tuple		,     __a				:	Any		,     __a				:	List[str]							)			->  Optional[int]:


												"""simple docstring"""

												__lowercase      :							int				=     self.num_labels
												__lowercase      :							List[Any]				=     MobileNetVaForSemanticSegmentation(__a							)
												model.to(__a							)
												model.eval()
												__lowercase      :							Dict				=     model(__a							)
												self.parent.assertEqual(
												    result.logits.shape		,     (
												        self.batch_size,
												        self.num_labels,
												        self.image_size // self.output_stride,
												        self.image_size // self.output_stride,
												    )		,     )
												__lowercase      :							str				=     model(__a		,     labels=__a							)
												self.parent.assertEqual(
												    result.logits.shape		,     (
												        self.batch_size,
												        self.num_labels,
												        self.image_size // self.output_stride,
												        self.image_size // self.output_stride,
												    )		,     )



					def       lowerCAmelCase				(  self				:	Tuple							)			->  Optional[int]:


												"""simple docstring"""

												__lowercase      :							List[str]				=     self.prepare_config_and_inputs()
												__lowercase			,			__lowercase			,			__lowercase			,			__lowercase      :							List[str]				=     config_and_inputs
												__lowercase      :							List[Any]				=     {"""pixel_values""": pixel_values}
												return config, inputs_dict







@require_torch
class 		lowerCAmelCase     ( __a      ,       __a      ,       unittest.TestCase			):

					'''simple docstring'''

					_A						:			Tuple									=      (
					    (MobileNetVaModel, MobileNetVaForImageClassification, MobileNetVaForSemanticSegmentation)
					    if is_torch_available()
					    else ()
					)
					_A						:			Optional[Any]									=      (
					    {
					        '''feature-extraction''': MobileNetVaModel,
					        '''image-classification''': MobileNetVaForImageClassification,
					        '''image-segmentation''': MobileNetVaForSemanticSegmentation,
					    }
					    if is_torch_available()
					    else {}
					)

					_A						:			Tuple									=      False
					_A						:			List[str]									=      False
					_A						:			List[str]									=      False
					_A						:			Optional[int]									=      False



					def       lowerCAmelCase				(  self				:	Optional[Any]							)			->  List[Any]:


												"""simple docstring"""

												__lowercase      :							Union[str, Any]				=     MobileNetVaModelTester(self							)
												__lowercase      :							int				=     MobileNetVaConfigTester(self		,     config_class=__a		,     has_text_modality=__a							)



					def       lowerCAmelCase				(  self				:	Union[str, Any]							)			->  List[Any]:


												"""simple docstring"""

												self.config_tester.run_common_tests()



					@unittest.skip(reason="""MobileNetV2 does not use inputs_embeds"""							)
					def       lowerCAmelCase				(  self				:	List[str]							)			->  Union[str, Any]:


												"""simple docstring"""

												pass



					@unittest.skip(reason="""MobileNetV2 does not support input and output embeddings"""							)
					def       lowerCAmelCase				(  self				:	Any							)			->  Tuple:


												"""simple docstring"""

												pass



					@unittest.skip(reason="""MobileNetV2 does not output attentions"""							)
					def       lowerCAmelCase				(  self				:	List[str]							)			->  int:


												"""simple docstring"""

												pass



					def       lowerCAmelCase				(  self				:	List[str]							)			->  Dict:


												"""simple docstring"""

												__lowercase			,			__lowercase      :							Optional[int]				=     self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																			__lowercase      :							List[Any]				=     model_class(__a							)
																			__lowercase      :							Any				=     inspect.signature(model.forward							)
																			# signature.parameters is an OrderedDict => so arg_names order is deterministic
																			__lowercase      :							int				=     [*signature.parameters.keys()]

																			__lowercase      :							Any				=     ["""pixel_values"""]
																			self.assertListEqual(arg_names[:1]		,     __a							)



					def       lowerCAmelCase				(  self				:	Dict							)			->  Any:


												"""simple docstring"""

												__lowercase      :							Union[str, Any]				=     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*__a							)



					def       lowerCAmelCase				(  self				:	List[str]							)			->  Tuple:


												"""simple docstring"""

												def check_hidden_states_output(__a				:	List[Any]		,     __a				:	Tuple		,     __a				:	List[str]							):
																			__lowercase      :							Optional[Any]				=     model_class(__a							)
																			model.to(__a							)
																			model.eval()

																			with torch.no_grad():
																										__lowercase      :							List[Any]				=     model(**self._prepare_for_class(__a		,     __a							)							)

																			__lowercase      :							Tuple				=     outputs.hidden_states

																			__lowercase      :							str				=     16
																			self.assertEqual(len(__a							)		,     __a							)

												__lowercase			,			__lowercase      :							int				=     self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																			__lowercase      :							Any				=     True
																			check_hidden_states_output(__a		,     __a		,     __a							)

																			# check that output_hidden_states also work using config
																			del inputs_dict["output_hidden_states"]
																			__lowercase      :							Union[str, Any]				=     True

																			check_hidden_states_output(__a		,     __a		,     __a							)



					def       lowerCAmelCase				(  self				:	Union[str, Any]							)			->  Any:


												"""simple docstring"""

												__lowercase      :							Dict				=     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_image_classification(*__a							)



					def       lowerCAmelCase				(  self				:	List[str]							)			->  int:


												"""simple docstring"""

												__lowercase      :							Union[str, Any]				=     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_semantic_segmentation(*__a							)



					@slow
					def       lowerCAmelCase				(  self				:	Optional[int]							)			->  Optional[Any]:


												"""simple docstring"""

												for model_name in MOBILENET_V2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			__lowercase      :							Optional[int]				=     MobileNetVaModel.from_pretrained(__a							)
																			self.assertIsNotNone(__a							)







def       snake_case_		(     ):
							__lowercase      :							List[Any]				=     Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""      )
							return image



@require_torch
@require_vision
class 		lowerCAmelCase     ( unittest.TestCase			):

					'''simple docstring'''

					@cached_property
					def       lowerCAmelCase				(  self				:	Any							)			->  Union[str, Any]:


												"""simple docstring"""

												return (
												    MobileNetVaImageProcessor.from_pretrained("""google/mobilenet_v2_1.0_224"""							) if is_vision_available() else None
												)



					@slow
					def       lowerCAmelCase				(  self				:	str							)			->  int:


												"""simple docstring"""

												__lowercase      :							Tuple				=     MobileNetVaForImageClassification.from_pretrained("""google/mobilenet_v2_1.0_224"""							).to(__a							)

												__lowercase      :							str				=     self.default_image_processor
												__lowercase      :							Tuple				=     prepare_img()
												__lowercase      :							Tuple				=     image_processor(images=__a		,     return_tensors="""pt"""							).to(__a							)

												# forward pass
												with torch.no_grad():
																			__lowercase      :							str				=     model(**__a							)

												# verify the logits
												__lowercase      :							Union[str, Any]				=     torch.Size((1, 1001)							)
												self.assertEqual(outputs.logits.shape		,     __a							)

												__lowercase      :							str				=     torch.tensor([0.2445, -1.1993, 0.1905]							).to(__a							)

												self.assertTrue(torch.allclose(outputs.logits[0, :3]		,     __a		,     atol=1E-4							)							)



					@slow
					def       lowerCAmelCase				(  self				:	Tuple							)			->  Any:


												"""simple docstring"""

												__lowercase      :							int				=     MobileNetVaForSemanticSegmentation.from_pretrained("""google/deeplabv3_mobilenet_v2_1.0_513"""							)
												__lowercase      :							Dict				=     model.to(__a							)

												__lowercase      :							Tuple				=     MobileNetVaImageProcessor.from_pretrained("""google/deeplabv3_mobilenet_v2_1.0_513"""							)

												__lowercase      :							List[str]				=     prepare_img()
												__lowercase      :							Optional[int]				=     image_processor(images=__a		,     return_tensors="""pt"""							).to(__a							)

												# forward pass
												with torch.no_grad():
																			__lowercase      :							Union[str, Any]				=     model(**__a							)
												__lowercase      :							Any				=     outputs.logits

												# verify the logits
												__lowercase      :							Dict				=     torch.Size((1, 21, 65, 65)							)
												self.assertEqual(logits.shape		,     __a							)

												__lowercase      :							str				=     torch.tensor(
												    [
												        [[17.5790, 17.7581, 18.3355], [18.3257, 18.4230, 18.8973], [18.6169, 18.8650, 19.2187]],
												        [[-2.1595, -2.0977, -2.3741], [-2.4226, -2.3028, -2.6835], [-2.7819, -2.5991, -2.7706]],
												        [[4.2058, 4.8317, 4.7638], [4.4136, 5.0361, 4.9383], [4.5028, 4.9644, 4.8734]],
												    ]		,     device=__a		,     )

												self.assertTrue(torch.allclose(logits[0, :3, :3, :3]		,     __a		,     atol=1E-4							)							)