






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class 	_lowercase				(      A__						,      A__						,      unittest.TestCase	):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__    :      int      =		StableDiffusionParadigmsPipeline
				SCREAMING_SNAKE_CASE__    :      str      =		TEXT_TO_IMAGE_PARAMS
				SCREAMING_SNAKE_CASE__    :      Tuple      =		TEXT_TO_IMAGE_BATCH_PARAMS
				SCREAMING_SNAKE_CASE__    :      int      =		TEXT_TO_IMAGE_IMAGE_PARAMS
				SCREAMING_SNAKE_CASE__    :      Any      =		TEXT_TO_IMAGE_IMAGE_PARAMS


				def 		__magic_name__(	self      :Union[str, Any]     )		->      List[str]:
										torch.manual_seed(0     )
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		UNetaDConditionModel(
										    block_out_channels=(32, 64)      ,      layers_per_block=2      ,      sample_size=32      ,      in_channels=4      ,      out_channels=4      ,      down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')      ,      up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')      ,      cross_attention_dim=32      ,      attention_head_dim=(2, 4)      ,      use_linear_projection=lowerCAmelCase__      ,      )
										__SCREAMING_SNAKE_CASE						:					List[str]		      =		DDIMParallelScheduler(
										    beta_start=0.0_0085      ,      beta_end=0.012      ,      beta_schedule='''scaled_linear'''      ,      clip_sample=lowerCAmelCase__      ,      set_alpha_to_one=lowerCAmelCase__      ,      )
										torch.manual_seed(0     )
										__SCREAMING_SNAKE_CASE						:					Dict		      =		AutoencoderKL(
										    block_out_channels=[32, 64]      ,      in_channels=3      ,      out_channels=3      ,      down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']      ,      up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']      ,      latent_channels=4      ,      sample_size=128      ,      )
										torch.manual_seed(0     )
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		CLIPTextConfig(
										    bos_token_id=0      ,      eos_token_id=2      ,      hidden_size=32      ,      intermediate_size=37      ,      layer_norm_eps=1E-0_5      ,      num_attention_heads=4      ,      num_hidden_layers=5      ,      pad_token_id=1      ,      vocab_size=1_000      ,      hidden_act='''gelu'''      ,      projection_dim=512      ,      )
										__SCREAMING_SNAKE_CASE						:					int		      =		CLIPTextModel(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					str		      =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

										__SCREAMING_SNAKE_CASE						:					int		      =		{
										    '''unet''': unet,
										    '''scheduler''': scheduler,
										    '''vae''': vae,
										    '''text_encoder''': text_encoder,
										    '''tokenizer''': tokenizer,
										    '''safety_checker''': None,
										    '''feature_extractor''': None,
										}
										return components


				def 		__magic_name__(	self      :Optional[int]      ,      lowerCAmelCase__      :Tuple      ,      lowerCAmelCase__      :str=0     )		->      Union[str, Any]:
										if str(lowerCAmelCase__     ).startswith('''mps'''     ):
																__SCREAMING_SNAKE_CASE						:					Any		      =		torch.manual_seed(lowerCAmelCase__     )
										else:
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		torch.Generator(device=lowerCAmelCase__     ).manual_seed(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Any		      =		{
										    '''prompt''': '''a photograph of an astronaut riding a horse''',
										    '''generator''': generator,
										    '''num_inference_steps''': 10,
										    '''guidance_scale''': 6.0,
										    '''output_type''': '''numpy''',
										    '''parallel''': 3,
										    '''debug''': True,
										}
										return inputs


				def 		__magic_name__(	self      :List[str]     )		->      Tuple:
										__SCREAMING_SNAKE_CASE						:					List[str]		      =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		self.get_dummy_components()
										__SCREAMING_SNAKE_CASE						:					Any		      =		StableDiffusionParadigmsPipeline(**lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		sd_pipe.to(lowerCAmelCase__     )
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

										__SCREAMING_SNAKE_CASE						:					Dict		      =		self.get_dummy_inputs(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		sd_pipe(**lowerCAmelCase__     ).images
										__SCREAMING_SNAKE_CASE						:					Dict		      =		image[0, -3:, -3:, -1]
										assert image.shape == (1, 64, 64, 3)

										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]     )

										assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-2


				def 		__magic_name__(	self      :Any     )		->      Dict:
										__SCREAMING_SNAKE_CASE						:					int		      =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		self.get_dummy_components()
										torch.manual_seed(0     )
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		DDPMParallelScheduler()
										torch.manual_seed(0     )
										__SCREAMING_SNAKE_CASE						:					List[str]		      =		StableDiffusionParadigmsPipeline(**lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		sd_pipe.to(lowerCAmelCase__     )
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

										__SCREAMING_SNAKE_CASE						:					Dict		      =		self.get_dummy_inputs(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					List[str]		      =		sd_pipe(**lowerCAmelCase__     ).images
										__SCREAMING_SNAKE_CASE						:					Dict		      =		image[0, -3:, -3:, -1]
										assert image.shape == (1, 64, 64, 3)

										__SCREAMING_SNAKE_CASE						:					List[str]		      =		np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]     )

										assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-2


				def 		__magic_name__(	self      :int     )		->      Any:
										super().test_inference_batch_consistent(batch_sizes=[1, 2]     )


				def 		__magic_name__(	self      :Tuple     )		->      Optional[Any]:
										super().test_inference_batch_single_identical(batch_size=2      ,      expected_max_diff=3E-3     )



				def 		__magic_name__(	self      :int     )		->      Dict:
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE						:					Dict		      =		self.get_dummy_components()
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		StableDiffusionParadigmsPipeline(**lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					int		      =		sd_pipe.to(lowerCAmelCase__     )
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__     )

										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		self.get_dummy_inputs(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					str		      =		'''french fries'''
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		sd_pipe(**lowerCAmelCase__      ,      negative_prompt=lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		output.images
										__SCREAMING_SNAKE_CASE						:					Any		      =		image[0, -3:, -3:, -1]

										assert image.shape == (1, 64, 64, 3)

										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]     )

										assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-2




@slow
@require_torch_gpu
class 	_lowercase				(      unittest.TestCase	):






				'''simple docstring'''



				def 		__magic_name__(	self      :str     )		->      List[str]:
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


				def 		__magic_name__(	self      :str      ,      lowerCAmelCase__      :Dict=0     )		->      Any:
										__SCREAMING_SNAKE_CASE						:					str		      =		torch.Generator(device=lowerCAmelCase__     ).manual_seed(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		{
										    '''prompt''': '''a photograph of an astronaut riding a horse''',
										    '''generator''': generator,
										    '''num_inference_steps''': 10,
										    '''guidance_scale''': 7.5,
										    '''output_type''': '''numpy''',
										    '''parallel''': 3,
										    '''debug''': True,
										}
										return inputs



				def 		__magic_name__(	self      :Optional[int]     )		->      Any:
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		'''stabilityai/stable-diffusion-2-base'''
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		DDIMParallelScheduler.from_pretrained(lowerCAmelCase__      ,      subfolder='''scheduler'''     )
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		StableDiffusionParadigmsPipeline.from_pretrained(lowerCAmelCase__      ,      scheduler=lowerCAmelCase__      ,      safety_checker=lowerCAmelCase__     )
										pipe.to(lowerCAmelCase__     )
										pipe.set_progress_bar_config(disable=lowerCAmelCase__     )
										pipe.enable_attention_slicing()

										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		self.get_inputs()
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		pipe(**lowerCAmelCase__     ).images
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		image[0, -3:, -3:, -1].flatten()

										assert image.shape == (1, 512, 512, 3)

										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]     )

										assert np.abs(expected_slice - image_slice     ).max() < 1E-2




import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.albert.modeling_albert import (
    ALBERT_INPUTS_DOCSTRING,
    ALBERT_START_DOCSTRING,
    AlbertModel,
    AlbertPreTrainedModel,
    AlbertTransformer,
)


__lowerCAmelCase     :			Any      =logging.getLogger(__name__)
class 	_lowercase				(      A__	):






				'''simple docstring'''



				def 		__magic_name__(	self      :Optional[int]      ,      lowerCAmelCase__      :Dict      ,      lowerCAmelCase__      :List[str]      ,      lowerCAmelCase__      :List[Any]=None      ,      lowerCAmelCase__      :Optional[Any]=None     )		->      List[Any]:
										if current_layer == 0:
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		self.embedding_hidden_mapping_in(lowerCAmelCase__     )
										else:
																__SCREAMING_SNAKE_CASE						:					Tuple		      =		hidden_states[0]

										__SCREAMING_SNAKE_CASE						:					Tuple		      =		int(self.config.num_hidden_layers / self.config.num_hidden_groups     )

										# Index of the hidden group
										__SCREAMING_SNAKE_CASE						:					List[str]		      =		int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)     )

										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		self.albert_layer_groups[group_idx](
										    lowerCAmelCase__      ,      lowerCAmelCase__      ,      head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]      ,      )
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		layer_group_output[0]

										return (hidden_states,)




@add_start_docstrings(
    '''The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top.'''						,      A__						,      )
class 	_lowercase				(      A__	):






				'''simple docstring'''



				def __init__(	self      :int      ,      lowerCAmelCase__      :List[str]     )		->      List[Any]:
										super().__init__(lowerCAmelCase__     )

										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		AlbertTransformerWithPabee(lowerCAmelCase__     )

										self.init_weights()
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		0
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		0
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		0

										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		0


				def 		__magic_name__(	self      :int      ,      lowerCAmelCase__      :Optional[int]     )		->      str:
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		threshold


				def 		__magic_name__(	self      :List[Any]      ,      lowerCAmelCase__      :Optional[int]     )		->      List[Any]:
										__SCREAMING_SNAKE_CASE						:					Dict		      =		patience


				def 		__magic_name__(	self      :Union[str, Any]     )		->      str:
										__SCREAMING_SNAKE_CASE						:					Any		      =		0
										__SCREAMING_SNAKE_CASE						:					int		      =		0


				def 		__magic_name__(	self      :Optional[int]     )		->      Dict:
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		self.inference_layers_num / self.inference_instances_num
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		(
										    f'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
										    f''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
										)
										print(lowerCAmelCase__     )



				@add_start_docstrings_to_model_forward(lowerCAmelCase__     )
				def 		__magic_name__(	self      :Dict      ,      lowerCAmelCase__      :List[str]=None      ,      lowerCAmelCase__      :List[str]=None      ,      lowerCAmelCase__      :Optional[Any]=None      ,      lowerCAmelCase__      :Any=None      ,      lowerCAmelCase__      :Dict=None      ,      lowerCAmelCase__      :Optional[int]=None      ,      lowerCAmelCase__      :str=None      ,      lowerCAmelCase__      :str=None      ,      lowerCAmelCase__      :str=False      ,      )		->      Union[str, Any]:

										if input_ids is not None and inputs_embeds is not None:
																raise ValueError('''You cannot specify both input_ids and inputs_embeds at the same time'''     )
										elif input_ids is not None:
																__SCREAMING_SNAKE_CASE						:					List[Any]		      =		input_ids.size()
										elif inputs_embeds is not None:
																__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		inputs_embeds.size()[:-1]
										else:
																raise ValueError('''You have to specify either input_ids or inputs_embeds'''     )

										__SCREAMING_SNAKE_CASE						:					List[str]		      =		input_ids.device if input_ids is not None else inputs_embeds.device

										if attention_mask is None:
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		torch.ones(lowerCAmelCase__      ,      device=lowerCAmelCase__     )
										if token_type_ids is None:
																__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		torch.zeros(lowerCAmelCase__      ,      dtype=torch.long      ,      device=lowerCAmelCase__     )

										__SCREAMING_SNAKE_CASE						:					Any		      =		attention_mask.unsqueeze(1     ).unsqueeze(2     )
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		extended_attention_mask.to(dtype=self.dtype     )  # fp16 compatibility
										__SCREAMING_SNAKE_CASE						:					Dict		      =		(1.0 - extended_attention_mask) * -1_0000.0
										__SCREAMING_SNAKE_CASE						:					str		      =		self.get_head_mask(lowerCAmelCase__      ,      self.config.num_hidden_layers     )

										__SCREAMING_SNAKE_CASE						:					List[str]		      =		self.embeddings(
										    lowerCAmelCase__      ,      position_ids=lowerCAmelCase__      ,      token_type_ids=lowerCAmelCase__      ,      inputs_embeds=lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		embedding_output

										if self.training:
																__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		[]
																for i in range(self.config.num_hidden_layers     ):
																						__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		self.encoder.adaptive_forward(
																						    lowerCAmelCase__      ,      current_layer=lowerCAmelCase__      ,      attention_mask=lowerCAmelCase__      ,      head_mask=lowerCAmelCase__      ,      )

																						__SCREAMING_SNAKE_CASE						:					Any		      =		self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]     )     )
																						__SCREAMING_SNAKE_CASE						:					Tuple		      =		output_layers[i](output_dropout(lowerCAmelCase__     )     )
																						res.append(lowerCAmelCase__     )
										elif self.patience == 0:  # Use all layers for inference
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		self.encoder(lowerCAmelCase__      ,      lowerCAmelCase__      ,      head_mask=lowerCAmelCase__     )
																__SCREAMING_SNAKE_CASE						:					List[Any]		      =		self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]     )     )
																__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		[output_layers[self.config.num_hidden_layers - 1](lowerCAmelCase__     )]
										else:
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		0
																__SCREAMING_SNAKE_CASE						:					Any		      =		None
																__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		0
																for i in range(self.config.num_hidden_layers     ):
																						calculated_layer_num += 1
																						__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		self.encoder.adaptive_forward(
																						    lowerCAmelCase__      ,      current_layer=lowerCAmelCase__      ,      attention_mask=lowerCAmelCase__      ,      head_mask=lowerCAmelCase__      ,      )

																						__SCREAMING_SNAKE_CASE						:					List[str]		      =		self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]     )     )
																						__SCREAMING_SNAKE_CASE						:					str		      =		output_layers[i](lowerCAmelCase__     )
																						if regression:
																												__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		logits.detach()
																												if patient_result is not None:
																																		__SCREAMING_SNAKE_CASE						:					int		      =		patient_result.detach()
																												if (patient_result is not None) and torch.abs(patient_result - labels     ) < self.regression_threshold:
																																		patient_counter += 1
																												else:
																																		__SCREAMING_SNAKE_CASE						:					List[str]		      =		0
																						else:
																												__SCREAMING_SNAKE_CASE						:					List[str]		      =		logits.detach().argmax(dim=1     )
																												if patient_result is not None:
																																		__SCREAMING_SNAKE_CASE						:					List[str]		      =		patient_result.detach().argmax(dim=1     )
																												if (patient_result is not None) and torch.all(labels.eq(lowerCAmelCase__     )     ):
																																		patient_counter += 1
																												else:
																																		__SCREAMING_SNAKE_CASE						:					str		      =		0

																						__SCREAMING_SNAKE_CASE						:					str		      =		logits
																						if patient_counter == self.patience:
																												break
																__SCREAMING_SNAKE_CASE						:					List[str]		      =		[patient_result]
																self.inference_layers_num += calculated_layer_num
																self.inference_instances_num += 1

										return res




@add_start_docstrings(
    '''Albert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. '''						,      A__						,      )
class 	_lowercase				(      A__	):






				'''simple docstring'''



				def __init__(	self      :List[Any]      ,      lowerCAmelCase__      :int     )		->      Optional[int]:
										super().__init__(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Any		      =		config.num_labels

										__SCREAMING_SNAKE_CASE						:					str		      =		AlbertModelWithPabee(lowerCAmelCase__     )
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		nn.Dropout(config.classifier_dropout_prob     )
										__SCREAMING_SNAKE_CASE						:					Dict		      =		nn.ModuleList(
										    [nn.Linear(config.hidden_size      ,      self.config.num_labels     ) for _ in range(config.num_hidden_layers     )]     )

										self.init_weights()



				@add_start_docstrings_to_model_forward(lowerCAmelCase__     )
				def 		__magic_name__(	self      :Optional[int]      ,      lowerCAmelCase__      :str=None      ,      lowerCAmelCase__      :List[str]=None      ,      lowerCAmelCase__      :Dict=None      ,      lowerCAmelCase__      :str=None      ,      lowerCAmelCase__      :Tuple=None      ,      lowerCAmelCase__      :int=None      ,      lowerCAmelCase__      :Any=None      ,      )		->      Dict:
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		self.albert(
										    input_ids=lowerCAmelCase__      ,      attention_mask=lowerCAmelCase__      ,      token_type_ids=lowerCAmelCase__      ,      position_ids=lowerCAmelCase__      ,      head_mask=lowerCAmelCase__      ,      inputs_embeds=lowerCAmelCase__      ,      output_dropout=self.dropout      ,      output_layers=self.classifiers      ,      regression=self.num_labels == 1      ,      )

										__SCREAMING_SNAKE_CASE						:					int		      =		(logits[-1],)

										if labels is not None:
																__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		None
																__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		0
																for ix, logits_item in enumerate(lowerCAmelCase__     ):
																						if self.num_labels == 1:
																												#  We are doing regression
																												__SCREAMING_SNAKE_CASE						:					str		      =		MSELoss()
																												__SCREAMING_SNAKE_CASE						:					List[Any]		      =		loss_fct(logits_item.view(-1     )      ,      labels.view(-1     )     )
																						else:
																												__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		CrossEntropyLoss()
																												__SCREAMING_SNAKE_CASE						:					List[Any]		      =		loss_fct(logits_item.view(-1      ,      self.num_labels     )      ,      labels.view(-1     )     )
																						if total_loss is None:
																												__SCREAMING_SNAKE_CASE						:					Dict		      =		loss
																						else:
																												total_loss += loss * (ix + 1)
																						total_weights += ix + 1
																__SCREAMING_SNAKE_CASE						:					Any		      =		(total_loss / total_weights,) + outputs

										return outputs
