
from torch import nn

class        _UpperCamelCase				(					nn.Module  ):






			'''simple docstring'''

			def __init__(       self :       Dict						,	snake_case_ :       Tuple						,	snake_case_ :       Optional[int]      ):
							super().__init__()
							UpperCamelCase_:							Dict	=						class_size
							UpperCamelCase_:							Tuple	=						embed_size
							# self.mlp1 = nn.Linear(embed_size, embed_size)
							# self.mlp2 = (nn.Linear(embed_size, class_size))
							UpperCamelCase_:							Dict	=						nn.Linear(UpperCAmelCase_						,	UpperCAmelCase_      )




			def       lowerCAmelCase__	(       self :       Tuple						,	snake_case_ :       Union[str, Any]      ):
							# hidden_state = nn.functional.relu(self.mlp1(hidden_state))
							# hidden_state = self.mlp2(hidden_state)
							UpperCamelCase_:							int	=						self.mlp(UpperCAmelCase_      )
							return logits








import os

from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE, hf_cache_home


lowerCamelCase_						: Optional[int]         =  HUGGINGFACE_HUB_CACHE


lowerCamelCase_						: List[str]         =  """config.json"""
lowerCamelCase_						: Any         =  """diffusion_pytorch_model.bin"""
lowerCamelCase_						: Union[str, Any]         =  """diffusion_flax_model.msgpack"""
lowerCamelCase_						: Dict         =  """model.onnx"""
lowerCamelCase_						: List[Any]         =  """diffusion_pytorch_model.safetensors"""
lowerCamelCase_						: Optional[Any]         =  """weights.pb"""
lowerCamelCase_						: Optional[Any]         =  """https://huggingface.co"""
lowerCamelCase_						: Union[str, Any]         =  default_cache_path
lowerCamelCase_						: Tuple         =  """diffusers_modules"""
lowerCamelCase_						: Optional[Any]         =  os.getenv("""HF_MODULES_CACHE""", os.path.join(hf_cache_home, """modules"""))
lowerCamelCase_						: str         =  ["""fp16""", """non-ema"""]
lowerCamelCase_						: List[Any]         =  """.self_attn"""






