

import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, slow, torch_device

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin



class 					lowercase		( _UpperCamelCase				,       _UpperCamelCase				,       unittest.TestCase      ):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE      							=       UNetaDModel
 __SCREAMING_SNAKE_CASE      							=       """sample"""




 @property
 def 						_A							(self				)				->      List[str]:






  """simple docstring"""





  UpperCAmelCase__           =       4
  UpperCAmelCase__           =       14
  UpperCAmelCase__           =       16

  UpperCAmelCase__           =       floats_tensor((batch_size, num_features, seq_len)				).to(__a				)
  UpperCAmelCase__           =       torch.tensor([10] * batch_size				).to(__a				)

  return {"sample": noise, "timestep": time_step}




 @property
 def 						_A							(self				)				->      Union[str, Any]:






  """simple docstring"""





  return (4, 14, 16)




 @property
 def 						_A							(self				)				->      Dict:






  """simple docstring"""





  return (4, 14, 16)




 def 						_A							(self				)				->      List[str]:






  """simple docstring"""





  pass




 def 						_A							(self				)				->      Any:






  """simple docstring"""





  pass




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      List[str]:






  """simple docstring"""





  super().test_determinism()




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      Tuple:






  """simple docstring"""





  super().test_outputs_equivalence()




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      List[str]:






  """simple docstring"""





  super().test_from_save_pretrained()




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      Optional[int]:






  """simple docstring"""





  super().test_from_save_pretrained_variant()




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      Optional[int]:






  """simple docstring"""





  super().test_model_from_pretrained()




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      int:






  """simple docstring"""





  super().test_output()




 def 						_A							(self				)				->      List[str]:






  """simple docstring"""





  UpperCAmelCase__           =       {
      'block_out_channels': (32, 64, 128, 256),
      'in_channels': 14,
      'out_channels': 14,
      'time_embedding_type': 'positional',
      'use_timestep_embedding': True,
      'flip_sin_to_cos': False,
      'freq_shift': 1.0,
      'out_block_type': 'OutConv1DBlock',
      'mid_block_type': 'MidResTemporalBlock1D',
      'down_block_types': ('DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D'),
      'up_block_types': ('UpResnetBlock1D', 'UpResnetBlock1D', 'UpResnetBlock1D'),
      'act_fn': 'mish',
  }
  UpperCAmelCase__           =       self.dummy_input
  return init_dict, inputs_dict




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      Dict:






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       UNetaDModel.from_pretrained(
      'bglick13/hopper-medium-v2-value-function-hor32' , output_loading_info=__a , subfolder='unet'				)
  self.assertIsNotNone(__a				)
  self.assertEqual(len(loading_info['missing_keys']				) , 0				)

  model.to(__a				)
  UpperCAmelCase__           =       model(**self.dummy_input				)

  assert image is not None, "Make sure output is not None"




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      int:






  """simple docstring"""





  UpperCAmelCase__           =       UNetaDModel.from_pretrained('bglick13/hopper-medium-v2-value-function-hor32' , subfolder='unet'				)
  torch.manual_seed(0				)
  if torch.cuda.is_available():
   torch.cuda.manual_seed_all(0				)

  UpperCAmelCase__           =       model.config.in_channels
  UpperCAmelCase__           =       16
  UpperCAmelCase__           =       torch.randn((1, seq_len, num_features)				).permute(
      0 , 2 , 1				)  # match original, we can update values and remove
  UpperCAmelCase__           =       torch.full((num_features,) , 0				)

  with torch.no_grad():
   UpperCAmelCase__           =       model(__a , __a				).sample.permute(0 , 2 , 1				)

  UpperCAmelCase__           =       output[0, -3:, -3:].flatten()
  # fmt: off
  UpperCAmelCase__           =       torch.tensor([-2.13_71_72, 1.1_42_60_16, 0.3_68_86_87, -0.76_69_22, 0.7_30_31_46, 0.11_03_88_64, -0.4_76_06_33, 0.13_27_01_72, 0.02_59_13_48]				)
  # fmt: on
  self.assertTrue(torch.allclose(__a , __a , rtol=1E-3				)				)




 def 						_A							(self				)				->      List[Any]:






  """simple docstring"""





  pass




 @slow
 def 						_A							(self				)				->      Tuple:






  """simple docstring"""





  UpperCAmelCase__           =       'harmonai/maestro-150k'
  UpperCAmelCase__           =       UNetaDModel.from_pretrained(__a , subfolder='unet'				)
  model.to(__a				)

  UpperCAmelCase__           =       65536
  UpperCAmelCase__           =       torch.sin(torch.arange(__a				)[None, None, :].repeat(1 , 2 , 1				)				).to(__a				)
  UpperCAmelCase__           =       torch.tensor([1]				).to(__a				)

  with torch.no_grad():
   UpperCAmelCase__           =       model(__a , __a				).sample

  UpperCAmelCase__           =       output.abs().sum()
  UpperCAmelCase__           =       output.abs().max()

  assert (output_sum - 224.0896).abs() < 0.5
  assert (output_max - 0.06_07).abs() < 4E-4







class 					lowercase		( _UpperCamelCase				,       _UpperCamelCase				,       unittest.TestCase      ):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE      							=       UNetaDModel
 __SCREAMING_SNAKE_CASE      							=       """sample"""




 @property
 def 						_A							(self				)				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__           =       4
  UpperCAmelCase__           =       14
  UpperCAmelCase__           =       16

  UpperCAmelCase__           =       floats_tensor((batch_size, num_features, seq_len)				).to(__a				)
  UpperCAmelCase__           =       torch.tensor([10] * batch_size				).to(__a				)

  return {"sample": noise, "timestep": time_step}




 @property
 def 						_A							(self				)				->      int:






  """simple docstring"""





  return (4, 14, 16)




 @property
 def 						_A							(self				)				->      Optional[Any]:






  """simple docstring"""





  return (4, 14, 1)




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      Tuple:






  """simple docstring"""





  super().test_determinism()




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      Any:






  """simple docstring"""





  super().test_outputs_equivalence()




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      Optional[Any]:






  """simple docstring"""





  super().test_from_save_pretrained()




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      int:






  """simple docstring"""





  super().test_from_save_pretrained_variant()




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      Tuple:






  """simple docstring"""





  super().test_model_from_pretrained()




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.prepare_init_args_and_inputs_for_common()
  UpperCAmelCase__           =       self.model_class(**__a				)
  model.to(__a				)
  model.eval()

  with torch.no_grad():
   UpperCAmelCase__           =       model(**__a				)

   if isinstance(__a , __a				):
    UpperCAmelCase__           =       output.sample

  self.assertIsNotNone(__a				)
  UpperCAmelCase__           =       torch.Size((inputs_dict['sample'].shape[0], 1)				)
  self.assertEqual(output.shape , __a , 'Input and output shapes do not match'				)




 def 						_A							(self				)				->      Optional[int]:






  """simple docstring"""





  pass




 def 						_A							(self				)				->      Any:






  """simple docstring"""





  pass




 def 						_A							(self				)				->      Any:






  """simple docstring"""





  UpperCAmelCase__           =       {
      'in_channels': 14,
      'out_channels': 14,
      'down_block_types': ['DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D'],
      'up_block_types': [],
      'out_block_type': 'ValueFunction',
      'mid_block_type': 'ValueFunctionMidBlock1D',
      'block_out_channels': [32, 64, 128, 256],
      'layers_per_block': 1,
      'downsample_each_block': True,
      'use_timestep_embedding': True,
      'freq_shift': 1.0,
      'flip_sin_to_cos': False,
      'time_embedding_type': 'positional',
      'act_fn': 'mish',
  }
  UpperCAmelCase__           =       self.dummy_input
  return init_dict, inputs_dict




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      Union[str, Any]:






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       UNetaDModel.from_pretrained(
      'bglick13/hopper-medium-v2-value-function-hor32' , output_loading_info=__a , subfolder='value_function'				)
  self.assertIsNotNone(__a				)
  self.assertEqual(len(vf_loading_info['missing_keys']				) , 0				)

  value_function.to(__a				)
  UpperCAmelCase__           =       value_function(**self.dummy_input				)

  assert image is not None, "Make sure output is not None"




 @unittest.skipIf(torch_device == 'mps' , 'mish op not supported in MPS'				)
 def 						_A							(self				)				->      int:






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       UNetaDModel.from_pretrained(
      'bglick13/hopper-medium-v2-value-function-hor32' , output_loading_info=__a , subfolder='value_function'				)
  torch.manual_seed(0				)
  if torch.cuda.is_available():
   torch.cuda.manual_seed_all(0				)

  UpperCAmelCase__           =       value_function.config.in_channels
  UpperCAmelCase__           =       14
  UpperCAmelCase__           =       torch.randn((1, seq_len, num_features)				).permute(
      0 , 2 , 1				)  # match original, we can update values and remove
  UpperCAmelCase__           =       torch.full((num_features,) , 0				)

  with torch.no_grad():
   UpperCAmelCase__           =       value_function(__a , __a				).sample

  # fmt: off
  UpperCAmelCase__           =       torch.tensor([165.25] * seq_len				)
  # fmt: on
  self.assertTrue(torch.allclose(__a , __a , rtol=1E-3				)				)




 def 						_A							(self				)				->      Tuple:






  """simple docstring"""





  pass






import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VivitConfig
from transformers.models.auto import get_values
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


if is_torch_available():
 import torch
 from torch import nn

 from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VivitForVideoClassification, VivitModel
 from transformers.models.vivit.modeling_vivit import VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from transformers import VivitImageProcessor



class 					lowercase		:



 '''simple docstring'''


 def __init__(self , __a , __a=2 , __a=True , __a=True , __a=10 , __a=10 , __a=8 , __a=[2, 4, 4] , __a=3 , __a=768 , __a=5 , __a=4 , __a=37 , __a="gelu_fast" , __a=0.0 , __a=0.0 , __a=0.02 , __a=1E-0_6 , __a=True , __a=None , )				->      List[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       parent
  UpperCAmelCase__           =       batch_size
  UpperCAmelCase__           =       is_training
  UpperCAmelCase__           =       use_labels
  UpperCAmelCase__           =       num_labels
  UpperCAmelCase__           =       image_size
  UpperCAmelCase__           =       num_frames
  UpperCAmelCase__           =       tubelet_size
  UpperCAmelCase__           =       num_channels
  UpperCAmelCase__           =       hidden_size
  UpperCAmelCase__           =       num_hidden_layers
  UpperCAmelCase__           =       num_attention_heads
  UpperCAmelCase__           =       intermediate_size
  UpperCAmelCase__           =       hidden_act
  UpperCAmelCase__           =       hidden_dropout_prob
  UpperCAmelCase__           =       attention_probs_dropout_prob
  UpperCAmelCase__           =       initializer_range
  UpperCAmelCase__           =       layer_norm_eps
  UpperCAmelCase__           =       qkv_bias
  UpperCAmelCase__           =       scope

  UpperCAmelCase__           =       (
      (self.image_size // self.tubelet_size[2])
      * (self.image_size // self.tubelet_size[1])
      * (self.num_frames // self.tubelet_size[0])
  ) + 1  # CLS token




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






  """simple docstring"""





  UpperCAmelCase__           =       floats_tensor(
      [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]				)

  UpperCAmelCase__           =       None
  if self.use_labels:
   UpperCAmelCase__           =       ids_tensor([self.batch_size] , self.num_labels				)

  UpperCAmelCase__           =       self.get_config()

  return config, pixel_values, labels




 def 						UpperCamelCase__							(self				)				->      Tuple:






  """simple docstring"""





  UpperCAmelCase__           =       VivitConfig(
      num_frames=self.num_frames , image_size=self.image_size , tubelet_size=self.tubelet_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 , layer_norm_eps=self.layer_norm_eps , qkv_bias=self.qkv_bias , )
  UpperCAmelCase__           =       self.num_labels
  return config




 def 						UpperCamelCase__							(self , __a , __a , __a				)				->      str:






  """simple docstring"""





  UpperCAmelCase__           =       VivitModel(config=__a				)
  model.to(__a				)
  model.eval()
  UpperCAmelCase__           =       model(__a				)
  self.parent.assertEqual(result.last_hidden_state.shape , (self.batch_size, self.seq_length, self.hidden_size)				)




 def 						UpperCamelCase__							(self , __a , __a , __a				)				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__           =       VivitForVideoClassification(__a				)
  model.to(__a				)
  model.eval()

  UpperCAmelCase__           =       model(__a				)

  # verify the logits shape
  UpperCAmelCase__           =       torch.Size((self.batch_size, self.num_labels)				)
  self.parent.assertEqual(result.logits.shape , __a				)




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






  """simple docstring"""





  UpperCAmelCase__           =       self.prepare_config_and_inputs()
  UpperCAmelCase__       ,					UpperCAmelCase__       ,					UpperCAmelCase__           =       config_and_inputs
  UpperCAmelCase__           =       {'pixel_values': pixel_values}
  return config, inputs_dict







@require_torch
class 					lowercase		( _UpperCamelCase				,       unittest.TestCase      ):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE      							=       (VivitModel, VivitForVideoClassification) if is_torch_available() else ()

 __SCREAMING_SNAKE_CASE      							=       False
 __SCREAMING_SNAKE_CASE      							=       False
 __SCREAMING_SNAKE_CASE      							=       False
 __SCREAMING_SNAKE_CASE      							=       False




 def 						UpperCamelCase__							(self				)				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__           =       VivitModelTester(self				)
  UpperCAmelCase__           =       ConfigTester(self , config_class=__a , has_text_modality=__a , hidden_size=37				)




 def 						UpperCamelCase__							(self , __a , __a , __a=False				)				->      int:






  """simple docstring"""





  UpperCAmelCase__           =       copy.deepcopy(__a				)

  if return_labels:
   if model_class in get_values(__a				):
    UpperCAmelCase__           =       torch.zeros(
        self.model_tester.batch_size , dtype=torch.long , device=__a				)

  return inputs_dict




 def 						UpperCamelCase__							(self				)				->      Tuple:






  """simple docstring"""





  self.config_tester.run_common_tests()




 @unittest.skip(reason='Vivit does not use inputs_embeds'				)
 def 						UpperCamelCase__							(self				)				->      str:






  """simple docstring"""





  pass




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






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs_for_common()

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




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






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs_for_common()

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

   UpperCAmelCase__           =       ['pixel_values', 'head_mask']
   self.assertListEqual(arg_names[:2] , __a				)




 def 						UpperCamelCase__							(self				)				->      List[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*__a				)




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






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_for_video_classification(*__a				)




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






  """simple docstring"""





  for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   UpperCAmelCase__           =       VivitModel.from_pretrained(__a				)
   self.assertIsNotNone(__a				)




 def 						UpperCamelCase__							(self				)				->      Any:






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs_for_common()
  UpperCAmelCase__           =       True

  for model_class in self.all_model_classes:
   UpperCAmelCase__           =       self.model_tester.seq_length

   UpperCAmelCase__           =       True
   UpperCAmelCase__           =       False
   UpperCAmelCase__           =       True
   UpperCAmelCase__           =       model_class(__a				)
   model.to(__a				)
   model.eval()
   with torch.no_grad():
    UpperCAmelCase__           =       model(**self._prepare_for_class(__a , __a				)				)
   UpperCAmelCase__           =       outputs.attentions
   self.assertEqual(len(__a				) , self.model_tester.num_hidden_layers				)

   # check that output_attentions also work using config
   del inputs_dict["output_attentions"]
   UpperCAmelCase__           =       True
   UpperCAmelCase__           =       model_class(__a				)
   model.to(__a				)
   model.eval()
   with torch.no_grad():
    UpperCAmelCase__           =       model(**self._prepare_for_class(__a , __a				)				)
   UpperCAmelCase__           =       outputs.attentions
   self.assertEqual(len(__a				) , self.model_tester.num_hidden_layers				)

   self.assertListEqual(
       list(attentions[0].shape[-3:]				) , [self.model_tester.num_attention_heads, seq_len, seq_len] , )
   UpperCAmelCase__           =       len(__a				)

   # Check attention is always last and order is fine
   UpperCAmelCase__           =       True
   UpperCAmelCase__           =       True
   UpperCAmelCase__           =       model_class(__a				)
   model.to(__a				)
   model.eval()
   with torch.no_grad():
    UpperCAmelCase__           =       model(**self._prepare_for_class(__a , __a				)				)

   self.assertEqual(out_len + 1 , len(__a				)				)

   UpperCAmelCase__           =       outputs.attentions

   self.assertEqual(len(__a				) , self.model_tester.num_hidden_layers				)

   self.assertListEqual(
       list(self_attentions[0].shape[-3:]				) , [self.model_tester.num_attention_heads, seq_len, seq_len] , )




 def 						UpperCamelCase__							(self				)				->      Optional[int]:






  """simple docstring"""





  def check_hidden_states_output(__a , __a , __a				):
   UpperCAmelCase__           =       model_class(__a				)
   model.to(__a				)
   model.eval()

   with torch.no_grad():
    UpperCAmelCase__           =       model(**self._prepare_for_class(__a , __a				)				)

   UpperCAmelCase__           =       outputs.hidden_states
   UpperCAmelCase__           =       self.model_tester.num_hidden_layers + 1
   self.assertEqual(len(__a				) , __a				)

   UpperCAmelCase__           =       self.model_tester.seq_length

   self.assertListEqual(
       list(hidden_states[0].shape[-2:]				) , [seq_length, self.model_tester.hidden_size] , )

  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs_for_common()

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

   # check that output_hidden_states also work using config
   del inputs_dict["output_hidden_states"]
   UpperCAmelCase__           =       True

   check_hidden_states_output(__a , __a , __a				)




 @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'				)
 def 						UpperCamelCase__							(self				)				->      List[str]:






  """simple docstring"""





  pass







def 		UpperCamelCase_(    ) ->				List[Any]:
 UpperCAmelCase__           =       hf_hub_download(
     repo_id='hf-internal-testing/spaghetti-video'						,					filename='eating_spaghetti_32_frames.npy'						,					repo_type='dataset'     )
 UpperCAmelCase__           =       np.load(snake_case__     )
 return list(snake_case__     )







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



 '''simple docstring'''


 @cached_property
 def 						UpperCamelCase__							(self				)				->      List[str]:






  """simple docstring"""





  return VivitImageProcessor() if is_vision_available() else None




 @slow
 def 						UpperCamelCase__							(self				)				->      str:






  """simple docstring"""





  UpperCAmelCase__           =       VivitForVideoClassification.from_pretrained('google/vivit-b-16x2-kinetics400'				).to(__a				)

  UpperCAmelCase__           =       self.default_image_processor
  UpperCAmelCase__           =       prepare_video()
  UpperCAmelCase__           =       image_processor(__a , return_tensors='pt'				).to(__a				)

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

  # verify the logits
  UpperCAmelCase__           =       torch.Size((1, 400)				)
  self.assertEqual(outputs.logits.shape , __a				)

  # taken from original model
  UpperCAmelCase__           =       torch.tensor([-1.05_43, 2.07_64, -0.21_04, 0.44_39, -0.96_58]				).to(__a				)

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



