


'''simple docstring'''

import argparse
import os
import re

import numpy as np
import PIL
import torch
from timm import create_model
from torch.optim.lr_scheduler import OneCycleLR
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor

from accelerate import Accelerator





def lowerCamelCase     (     lowerCAmelCase							:      Tuple						):





	"""simple docstring"""







	__magic_name__			:       Optional[Any]	      =	fname.split(os.path.sep						)[-1]
	return re.search(R'^(.*)_\d+\.jpg$'						,						__lowerCAmelCase						).groups()[0]
class _lowerCamelCase			(		__a				):



		'''simple docstring'''







		def __init__(   self						:		Optional[Any]    ,	_A						:		Any    ,	_A						:		Any=None    ,	_A						:		int=None     )    ->					List[Any]:
			__magic_name__			:       int	      =	file_names
			__magic_name__			:       List[str]	      =	image_transform
			__magic_name__			:       Union[str, Any]	      =	label_to_id







		def __len__(   self						:		List[str]     )    ->					List[str]:
			return len(self.file_names     )







		def __getitem__(   self						:		Union[str, Any]    ,	_A						:		int     )    ->					Tuple:
			__magic_name__			:       Any	      =	self.file_names[idx]
			__magic_name__			:       Union[str, Any]	      =	PIL.Image.open(_UpperCamelCase     )
			__magic_name__			:       Optional[int]	      =	raw_image.convert('RGB'     )
			if self.image_transform is not None:
				__magic_name__			:       Optional[Any]	      =	self.image_transform(_UpperCamelCase     )
			__magic_name__			:       Optional[Any]	      =	extract_label(_UpperCamelCase     )
			if self.label_to_id is not None:
				__magic_name__			:       str	      =	self.label_to_id[label]
			return {"image": image, "label": label}







def lowerCamelCase     (     lowerCAmelCase							:      str						,						lowerCAmelCase							:      Any						):





	"""simple docstring"""







	if args.with_tracking:
		__magic_name__			:       Optional[int]	      =	Accelerator(
		    cpu=args.cpu						,						mixed_precision=args.mixed_precision						,						log_with='all'						,						project_dir=args.project_dir						)
	else:
		__magic_name__			:       Any	      =	Accelerator(cpu=args.cpu						,						mixed_precision=args.mixed_precision						)

	# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
	__magic_name__			:       Any	      =	config['lr']
	__magic_name__			:       Any	      =	int(config['num_epochs']						)
	__magic_name__			:       Union[str, Any]	      =	int(config['seed']						)
	__magic_name__			:       List[str]	      =	int(config['batch_size']						)
	__magic_name__			:       Dict	      =	config['image_size']
	if not isinstance(__lowerCAmelCase						,						(list, tuple)						):
		__magic_name__			:       Dict	      =	(image_size, image_size)

	# Parse out whether we are saving every epoch or after a certain number of batches
	if hasattr(args.checkpointing_steps						,						'isdigit'						):
		if args.checkpointing_steps == "epoch":
			__magic_name__			:       Union[str, Any]	      =	args.checkpointing_steps
		elif args.checkpointing_steps.isdigit():
			__magic_name__			:       str	      =	int(args.checkpointing_steps						)
		else:
			raise ValueError(
			    f'Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.'						)
	else:
		__magic_name__			:       Optional[Any]	      =	None

	# We need to initialize the trackers we use, and also store our configuration
	if args.with_tracking:
		__magic_name__			:       Optional[int]	      =	os.path.split(__lowerCAmelCase						)[-1].split('.'						)[0]
		accelerator.init_trackers(__lowerCAmelCase						,						__lowerCAmelCase						)

	# Grab all the image filenames
	__magic_name__			:       List[Any]	      =	[os.path.join(args.data_dir						,						__lowerCAmelCase						) for fname in os.listdir(args.data_dir						) if fname.endswith('.jpg'						)]

	# Build the label correspondences
	__magic_name__			:       List[Any]	      =	[extract_label(__lowerCAmelCase						) for fname in file_names]
	__magic_name__			:       str	      =	list(set(__lowerCAmelCase						)						)
	id_to_label.sort()
	__magic_name__			:       Dict	      =	{lbl: i for i, lbl in enumerate(__lowerCAmelCase						)}

	# Set the seed before splitting the data.
	np.random.seed(__lowerCAmelCase						)
	torch.manual_seed(__lowerCAmelCase						)
	torch.cuda.manual_seed_all(__lowerCAmelCase						)

	# Split our filenames between train and validation
	__magic_name__			:       List[Any]	      =	np.random.permutation(len(__lowerCAmelCase						)						)
	__magic_name__			:       List[str]	      =	int(0.8 * len(__lowerCAmelCase						)						)
	__magic_name__			:       Union[str, Any]	      =	random_perm[:cut]
	__magic_name__			:       List[str]	      =	random_perm[cut:]

	# For training we use a simple RandomResizedCrop
	__magic_name__			:       Any	      =	Compose([RandomResizedCrop(__lowerCAmelCase						,						scale=(0.5, 1.0)						), ToTensor()]						)
	__magic_name__			:       Tuple	      =	PetsDataset(
	    [file_names[i] for i in train_split]						,						image_transform=__lowerCAmelCase						,						label_to_id=__lowerCAmelCase						)

	# For evaluation, we use a deterministic Resize
	__magic_name__			:       Optional[Any]	      =	Compose([Resize(__lowerCAmelCase						), ToTensor()]						)
	__magic_name__			:       int	      =	PetsDataset([file_names[i] for i in eval_split]						,						image_transform=__lowerCAmelCase						,						label_to_id=__lowerCAmelCase						)

	# Instantiate dataloaders.
	__magic_name__			:       Optional[int]	      =	DataLoader(__lowerCAmelCase						,						shuffle=__lowerCAmelCase						,						batch_size=__lowerCAmelCase						,						num_workers=4						)
	__magic_name__			:       Dict	      =	DataLoader(__lowerCAmelCase						,						shuffle=__lowerCAmelCase						,						batch_size=__lowerCAmelCase						,						num_workers=4						)

	# Instantiate the model (we build the model here so that the seed also control new weights initialization)
	__magic_name__			:       Dict	      =	create_model('resnet50d'						,						pretrained=__lowerCAmelCase						,						num_classes=len(__lowerCAmelCase						)						)

	# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
	# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
	# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
	__magic_name__			:       int	      =	model.to(accelerator.device						)

	# Freezing the base model
	for param in model.parameters():
		__magic_name__			:       str	      =	False
	for param in model.get_classifier().parameters():
		__magic_name__			:       Optional[Any]	      =	True

	# We normalize the batches of images to be a bit faster.
	__magic_name__			:       Optional[Any]	      =	torch.tensor(model.default_cfg['mean']						)[None, :, None, None].to(accelerator.device						)
	__magic_name__			:       Union[str, Any]	      =	torch.tensor(model.default_cfg['std']						)[None, :, None, None].to(accelerator.device						)

	# Instantiate optimizer
	__magic_name__			:       Tuple	      =	torch.optim.Adam(params=model.parameters()						,						lr=lr / 25						)

	# Instantiate learning rate scheduler
	__magic_name__			:       Any	      =	OneCycleLR(optimizer=__lowerCAmelCase						,						max_lr=__lowerCAmelCase						,						epochs=__lowerCAmelCase						,						steps_per_epoch=len(__lowerCAmelCase						)						)

	# Prepare everything
	# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
	# prepare method.
	__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__      ,							__magic_name__			:       List[str]	      =	accelerator.prepare(
	    __lowerCAmelCase						,						__lowerCAmelCase						,						__lowerCAmelCase						,						__lowerCAmelCase						,						__lowerCAmelCase						)
	# We need to keep track of how many total steps we have iterated over
	__magic_name__			:       str	      =	0
	# We also need to keep track of the starting epoch so files are named properly
	__magic_name__			:       Optional[Any]	      =	0

	# Potentially load in the weights and states from a previous save
	if args.resume_from_checkpoint:
		if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
			accelerator.print(f'Resumed from checkpoint: {args.resume_from_checkpoint}'						)
			accelerator.load_state(args.resume_from_checkpoint						)
			__magic_name__			:       Optional[int]	      =	os.path.basename(args.resume_from_checkpoint						)
		else:
			# Get the most recent checkpoint
			__magic_name__			:       Optional[int]	      =	[f.name for f in os.scandir(os.getcwd()						) if f.is_dir()]
			dirs.sort(key=os.path.getctime						)
			__magic_name__			:       Union[str, Any]	      =	dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
		# Extract `epoch_{i}` or `step_{i}`
		__magic_name__			:       Tuple	      =	os.path.splitext(__lowerCAmelCase						)[0]

		if "epoch" in training_difference:
			__magic_name__			:       Optional[Any]	      =	int(training_difference.replace('epoch_'						,						''						)						) + 1
			__magic_name__			:       Dict	      =	None
		else:
			__magic_name__			:       Optional[Any]	      =	int(training_difference.replace('step_'						,						''						)						)
			__magic_name__			:       List[Any]	      =	resume_step // len(__lowerCAmelCase						)
			resume_step -= starting_epoch * len(__lowerCAmelCase						)

    # Now we train the model
	for epoch in range(__lowerCAmelCase						,						__lowerCAmelCase						):
		model.train()
		if args.with_tracking:
			__magic_name__			:       str	      =	0
		if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
			# We need to skip steps until we reach the resumed step
			__magic_name__			:       Dict	      =	accelerator.skip_first_batches(__lowerCAmelCase						,						__lowerCAmelCase						)
			overall_step += resume_step
		else:
			# After the first iteration though, we need to go back to the original dataloader
			__magic_name__			:       str	      =	train_dataloader
		for batch in active_dataloader:
			# We could avoid this line since we set the accelerator with `device_placement=True`.
			__magic_name__			:       Optional[int]	      =	{k: v.to(accelerator.device						) for k, v in batch.items()}
			__magic_name__			:       Union[str, Any]	      =	(batch['image'] - mean) / std
			__magic_name__			:       Union[str, Any]	      =	model(__lowerCAmelCase						)
			__magic_name__			:       List[str]	      =	torch.nn.functional.cross_entropy(__lowerCAmelCase						,						batch['label']						)
			# We keep track of the loss at each epoch
			if args.with_tracking:
				total_loss += loss.detach().float()
			accelerator.backward(__lowerCAmelCase						)
			optimizer.step()
			lr_scheduler.step()
			optimizer.zero_grad()
			overall_step += 1
			if isinstance(__lowerCAmelCase						,						__lowerCAmelCase						):
				__magic_name__			:       Optional[Any]	      =	f'step_{overall_step}'
				if overall_step % checkpointing_steps == 0:
					if args.output_dir is not None:
						__magic_name__			:       Optional[Any]	      =	os.path.join(args.output_dir						,						__lowerCAmelCase						)
					accelerator.save_state(__lowerCAmelCase						)
		model.eval()
		__magic_name__			:       Optional[int]	      =	0
		__magic_name__			:       Any	      =	0
		for step, batch in enumerate(__lowerCAmelCase						):
			# We could avoid this line since we set the accelerator with `device_placement=True`.
			__magic_name__			:       str	      =	{k: v.to(accelerator.device						) for k, v in batch.items()}
			__magic_name__			:       Optional[int]	      =	(batch['image'] - mean) / std
			with torch.no_grad():
				__magic_name__			:       Tuple	      =	model(__lowerCAmelCase						)
			__magic_name__			:       Optional[Any]	      =	outputs.argmax(dim=-1						)
			__magic_name__      ,							__magic_name__			:       List[str]	      =	accelerator.gather_for_metrics((predictions, batch['label'])						)
			__magic_name__			:       Optional[int]	      =	predictions == references
			num_elems += accurate_preds.shape[0]
			accurate += accurate_preds.long().sum()

		__magic_name__			:       List[Any]	      =	accurate.item() / num_elems
		# Use accelerator.print to print only on the main process.
		accelerator.print(f'epoch {epoch}: {100 * eval_metric:.2f}'						)
		if args.with_tracking:
			accelerator.log(
			    {
			        'accuracy': 100 * eval_metric,
			        'train_loss': total_loss.item() / len(__lowerCAmelCase						),
			        'epoch': epoch,
			    }						,						step=__lowerCAmelCase						,						)
		if checkpointing_steps == "epoch":
			__magic_name__			:       Dict	      =	f'epoch_{epoch}'
			if args.output_dir is not None:
				__magic_name__			:       Union[str, Any]	      =	os.path.join(args.output_dir						,						__lowerCAmelCase						)
			accelerator.save_state(__lowerCAmelCase						)

	if args.with_tracking:
		accelerator.end_training()





def lowerCamelCase     (     ):





	"""simple docstring"""







	__magic_name__			:       int	      =	argparse.ArgumentParser(description='Simple example of training script.'						)
	parser.add_argument('--data_dir'						,						required=__lowerCAmelCase						,						help='The data folder on disk.'						)
	parser.add_argument('--fp16'						,						action='store_true'						,						help='If passed, will use FP16 training.'						)
	parser.add_argument(
	    '--mixed_precision'						,						type=__lowerCAmelCase						,						default=__lowerCAmelCase						,						choices=['no', 'fp16', 'bf16', 'fp8']						,						help='Whether to use mixed precision. Choose'
	    'between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'
	    'and an Nvidia Ampere GPU.'						,						)
	parser.add_argument('--cpu'						,						action='store_true'						,						help='If passed, will train on the CPU.'						)
	parser.add_argument(
	    '--checkpointing_steps'						,						type=__lowerCAmelCase						,						default=__lowerCAmelCase						,						help='Whether the various states should be saved at the end of every n steps, or \'epoch\' for each epoch.'						,						)
	parser.add_argument(
	    '--output_dir'						,						type=__lowerCAmelCase						,						default='.'						,						help='Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'						,						)
	parser.add_argument(
	    '--resume_from_checkpoint'						,						type=__lowerCAmelCase						,						default=__lowerCAmelCase						,						help='If the training should continue from a checkpoint folder.'						,						)
	parser.add_argument(
	    '--with_tracking'						,						action='store_true'						,						help='Whether to load in all available experiment trackers from the environment and use them for logging.'						,						)
	parser.add_argument(
	    '--project_dir'						,						type=__lowerCAmelCase						,						default='logs'						,						help='Location on where to store experiment tracking logs` and relevent project information'						,						)
	__magic_name__			:       Any	      =	parser.parse_args()
	__magic_name__			:       int	      =	{'lr': 3e-2, 'num_epochs': 3, 'seed': 42, 'batch_size': 64, 'image_size': 224}
	training_function(__lowerCAmelCase						,						__lowerCAmelCase						)


if __name__ == "__main__":
							main()

import inspect
import unittest

from transformers import ViTMSNConfig
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 torch import nn

			from transformers import ViTMSNForImageClassification, ViTMSNModel
			from transformers.models.vit_msn.modeling_vit_msn import VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
			from PIL import Image

			from transformers import ViTImageProcessor





class       _lowerCAmelCase   :




				def __init__(   self							,							_UpperCamelCase							,							_UpperCamelCase=13							,							_UpperCamelCase=30							,							_UpperCamelCase=2							,							_UpperCamelCase=3							,							_UpperCamelCase=True							,							_UpperCamelCase=True							,							_UpperCamelCase=32							,							_UpperCamelCase=5							,							_UpperCamelCase=4							,							_UpperCamelCase=37							,							_UpperCamelCase="gelu"							,							_UpperCamelCase=0.1							,							_UpperCamelCase=0.1							,							_UpperCamelCase=10							,							_UpperCamelCase=0.02							,							_UpperCamelCase=None							,							)							-> List[Any]:
							lowerCAmelCase_         =      parent
							lowerCAmelCase_         =      batch_size
							lowerCAmelCase_         =      image_size
							lowerCAmelCase_         =      patch_size
							lowerCAmelCase_         =      num_channels
							lowerCAmelCase_         =      is_training
							lowerCAmelCase_         =      use_labels
							lowerCAmelCase_         =      hidden_size
							lowerCAmelCase_         =      num_hidden_layers
							lowerCAmelCase_         =      num_attention_heads
							lowerCAmelCase_         =      intermediate_size
							lowerCAmelCase_         =      hidden_act
							lowerCAmelCase_         =      hidden_dropout_prob
							lowerCAmelCase_         =      attention_probs_dropout_prob
							lowerCAmelCase_         =      type_sequence_label_size
							lowerCAmelCase_         =      initializer_range
							lowerCAmelCase_         =      scope

							# in ViT MSN, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
							lowerCAmelCase_         =      (image_size // patch_size) ** 2
							lowerCAmelCase_         =      num_patches + 1




				def 	__a		(   self	)							-> Union[str, Any]:
							lowerCAmelCase_         =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

							lowerCAmelCase_         =      None
							if self.use_labels:
										lowerCAmelCase_         =      ids_tensor([self.batch_size]							,							self.type_sequence_label_size	)

							lowerCAmelCase_         =      self.get_config()

							return config, pixel_values, labels




				def 	__a		(   self	)							-> Optional[Any]:
							return ViTMSNConfig(
							    image_size=self.image_size							,							patch_size=self.patch_size							,							num_channels=self.num_channels							,							hidden_size=self.hidden_size							,							num_hidden_layers=self.num_hidden_layers							,							num_attention_heads=self.num_attention_heads							,							intermediate_size=self.intermediate_size							,							hidden_act=self.hidden_act							,							hidden_dropout_prob=self.hidden_dropout_prob							,							attention_probs_dropout_prob=self.attention_probs_dropout_prob							,							initializer_range=self.initializer_range							,							)




				def 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase	)							-> str:
							lowerCAmelCase_         =      ViTMSNModel(config=_UpperCamelCase	)
							model.to(_UpperCamelCase	)
							model.eval()
							lowerCAmelCase_         =      model(_UpperCamelCase	)
							self.parent.assertEqual(result.last_hidden_state.shape							,							(self.batch_size, self.seq_length, self.hidden_size)	)




				def 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase	)							-> str:
							lowerCAmelCase_         =      self.type_sequence_label_size
							lowerCAmelCase_         =      ViTMSNForImageClassification(_UpperCamelCase	)
							model.to(_UpperCamelCase	)
							model.eval()
							lowerCAmelCase_         =      model(_UpperCamelCase							,							labels=_UpperCamelCase	)
							print("Pixel and labels shape: {pixel_values.shape}, {labels.shape}"	)
							print("Labels: {labels}"	)
							self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.type_sequence_label_size)	)

							# test greyscale images
							lowerCAmelCase_         =      1
							lowerCAmelCase_         =      ViTMSNForImageClassification(_UpperCamelCase	)
							model.to(_UpperCamelCase	)
							model.eval()

							lowerCAmelCase_         =      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
							lowerCAmelCase_         =      model(_UpperCamelCase	)
							self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.type_sequence_label_size)	)




				def 	__a		(   self	)							-> Union[str, Any]:
							lowerCAmelCase_         =      self.prepare_config_and_inputs()
							lowerCAmelCase_      ,    lowerCAmelCase_      ,    lowerCAmelCase_         =      config_and_inputs
							lowerCAmelCase_         =      {"pixel_values": pixel_values}
							return config, inputs_dict





@require_torch
class       _lowerCAmelCase   (      __a   , __a   , unittest.TestCase			):
				_lowercase     						=(ViTMSNModel, ViTMSNForImageClassification) if is_torch_available() else ()
				_lowercase     						=(
				    {'''feature-extraction''': ViTMSNModel, '''image-classification''': ViTMSNForImageClassification}
				    if is_torch_available()
				    else {}
				)

				_lowercase     						=False
				_lowercase     						=False
				_lowercase     						=False
				_lowercase     						=False




				def 	__a		(   self	)							-> Optional[int]:
							lowerCAmelCase_         =      ViTMSNModelTester(self	)
							lowerCAmelCase_         =      ConfigTester(self							,							config_class=_UpperCamelCase							,							has_text_modality=_UpperCamelCase							,							hidden_size=37	)




				def 	__a		(   self	)							-> Optional[int]:
							self.config_tester.run_common_tests()




				@unittest.skip(reason="ViTMSN does not use inputs_embeds"	)
				def 	__a		(   self	)							-> str:
							pass




				def 	__a		(   self	)							-> Any:
							lowerCAmelCase_      ,    lowerCAmelCase_         =      self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										lowerCAmelCase_         =      model_class(_UpperCamelCase	)
										self.assertIsInstance(model.get_input_embeddings()							,							(nn.Module)	)
										lowerCAmelCase_         =      model.get_output_embeddings()
										self.assertTrue(x is None or isinstance(_UpperCamelCase							,							nn.Linear	)	)




				def 	__a		(   self	)							-> Optional[Any]:
							lowerCAmelCase_      ,    lowerCAmelCase_         =      self.model_tester.prepare_config_and_inputs_for_common()

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

										lowerCAmelCase_         =      ["pixel_values"]
										self.assertListEqual(arg_names[:1]							,							_UpperCamelCase	)




				def 	__a		(   self	)							-> str:
							lowerCAmelCase_         =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_UpperCamelCase	)




				def 	__a		(   self	)							-> Any:
							lowerCAmelCase_         =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_UpperCamelCase	)




				@slow
				def 	__a		(   self	)							-> Union[str, Any]:
							for model_name in VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase_         =      ViTMSNModel.from_pretrained(_UpperCamelCase	)
										self.assertIsNotNone(_UpperCamelCase	)






def   lowerCamelCase__							(							):




			"""simple docstring"""


			lowerCAmelCase_         =      Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"							)
			return image





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




				@cached_property
				def 	__a		(   self	)							-> List[str]:
							return ViTImageProcessor.from_pretrained("facebook/vit-msn-small"	) if is_vision_available() else None




				@slow
				def 	__a		(   self	)							-> int:
							torch.manual_seed(2	)
							lowerCAmelCase_         =      ViTMSNForImageClassification.from_pretrained("facebook/vit-msn-small"	).to(_UpperCamelCase	)

							lowerCAmelCase_         =      self.default_image_processor
							lowerCAmelCase_         =      prepare_img()
							lowerCAmelCase_         =      image_processor(images=_UpperCamelCase							,							return_tensors="pt"	).to(_UpperCamelCase	)

							# forward pass
							with torch.no_grad():
										lowerCAmelCase_         =      model(**_UpperCamelCase	)

							# verify the logits
							lowerCAmelCase_         =      torch.Size((1, 1_000)	)
							self.assertEqual(outputs.logits.shape							,							_UpperCamelCase	)

							lowerCAmelCase_         =      torch.tensor([-0.0803, -0.4454, -0.2375]	).to(_UpperCamelCase	)

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



