




import gc
import unittest

from transformers import CTRLConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
							    CTRLForSequenceClassification,
							    CTRLLMHeadModel,
							    CTRLModel,
							)




class 		__A   :

						def __init__(		self :Tuple     ,			__snake_case :int     ,			__snake_case :int=14     ,			__snake_case :Tuple=7     ,			__snake_case :Any=True     ,			__snake_case :Dict=True     ,			__snake_case :Any=True     ,			__snake_case :Optional[Any]=True     ,			__snake_case :Optional[Any]=True     ,			__snake_case :Dict=99     ,			__snake_case :Any=32     ,			__snake_case :int=5     ,			__snake_case :Any=4     ,			__snake_case :str=37     ,			__snake_case :List[Any]="gelu"     ,			__snake_case :str=0.1     ,			__snake_case :str=0.1     ,			__snake_case :str=5_12     ,			__snake_case :int=16     ,			__snake_case :Optional[Any]=2     ,			__snake_case :Dict=0.02     ,			__snake_case :Tuple=3     ,			__snake_case :List[str]=4     ,			__snake_case :Any=None     ,			):
										'''simple docstring'''

										__magic_name__						:						int										=parent
										__magic_name__						:						str										=batch_size
										__magic_name__						:						str										=seq_length
										__magic_name__						:						Tuple										=is_training
										__magic_name__						:						List[str]										=use_token_type_ids
										__magic_name__						:						Tuple										=use_input_mask
										__magic_name__						:						Any										=use_labels
										__magic_name__						:						str										=use_mc_token_ids
										__magic_name__						:						List[Any]										=vocab_size
										__magic_name__						:						List[str]										=hidden_size
										__magic_name__						:						Tuple										=num_hidden_layers
										__magic_name__						:						Union[str, Any]										=num_attention_heads
										__magic_name__						:						Tuple										=intermediate_size
										__magic_name__						:						Optional[int]										=hidden_act
										__magic_name__						:						Tuple										=hidden_dropout_prob
										__magic_name__						:						str										=attention_probs_dropout_prob
										__magic_name__						:						Dict										=max_position_embeddings
										__magic_name__						:						Dict										=type_vocab_size
										__magic_name__						:						Any										=type_sequence_label_size
										__magic_name__						:						Optional[Any]										=initializer_range
										__magic_name__						:						Optional[Any]										=num_labels
										__magic_name__						:						Any										=num_choices
										__magic_name__						:						List[str]										=scope
										__magic_name__						:						Optional[Any]										=self.vocab_size - 1

						def      A__    (		self :List[Any]			):
										'''simple docstring'''

										__magic_name__						:						Tuple										=ids_tensor([self.batch_size, self.seq_length]     ,			self.vocab_size			)

										__magic_name__						:						Any										=None
										if self.use_input_mask:
														__magic_name__						:						Optional[int]										=random_attention_mask([self.batch_size, self.seq_length]			)

										__magic_name__						:						Any										=None
										if self.use_token_type_ids:
														__magic_name__						:						int										=ids_tensor([self.batch_size, self.seq_length]     ,			self.type_vocab_size			)

										__magic_name__						:						Tuple										=None
										if self.use_mc_token_ids:
														__magic_name__						:						Tuple										=ids_tensor([self.batch_size, self.num_choices]     ,			self.seq_length			)

										__magic_name__						:						Optional[Any]										=None
										__magic_name__						:						str										=None
										__magic_name__						:						List[str]										=None
										if self.use_labels:
														__magic_name__						:						int										=ids_tensor([self.batch_size]     ,			self.type_sequence_label_size			)
														__magic_name__						:						List[Any]										=ids_tensor([self.batch_size, self.seq_length]     ,			self.num_labels			)
														__magic_name__						:						List[Any]										=ids_tensor([self.batch_size]     ,			self.num_choices			)

										__magic_name__						:						Optional[int]										=self.get_config()

										__magic_name__						:						List[Any]										=ids_tensor([self.num_hidden_layers, self.num_attention_heads]     ,			2			)

										return (
										    config,
										    input_ids,
										    input_mask,
										    head_mask,
										    token_type_ids,
										    mc_token_ids,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										)

						def      A__    (		self :List[Any]			):
										'''simple docstring'''

										return CTRLConfig(
										    vocab_size=self.vocab_size     ,			n_embd=self.hidden_size     ,			n_layer=self.num_hidden_layers     ,			n_head=self.num_attention_heads     ,			n_positions=self.max_position_embeddings     ,			pad_token_id=self.pad_token_id     ,			)

						def      A__    (		self :int     ,			__snake_case :Dict     ,			__snake_case :Optional[Any]     ,			__snake_case :Tuple     ,			__snake_case :Dict     ,			__snake_case :Dict     ,			*__snake_case :str			):
										'''simple docstring'''

										__magic_name__						:						Optional[int]										=CTRLModel(config=__snake_case			)
										model.to(__snake_case			)
										model.eval()

										model(__snake_case     ,			token_type_ids=__snake_case     ,			head_mask=__snake_case			)
										model(__snake_case     ,			token_type_ids=__snake_case			)
										__magic_name__						:						int										=model(__snake_case			)
										self.parent.assertEqual(result.last_hidden_state.shape     ,			(self.batch_size, self.seq_length, self.hidden_size)			)
										self.parent.assertEqual(len(result.past_key_values			)     ,			config.n_layer			)

						def      A__    (		self :Dict     ,			__snake_case :Tuple     ,			__snake_case :str     ,			__snake_case :Dict     ,			__snake_case :Tuple     ,			__snake_case :List[Any]     ,			*__snake_case :List[str]			):
										'''simple docstring'''

										__magic_name__						:						str										=CTRLLMHeadModel(__snake_case			)
										model.to(__snake_case			)
										model.eval()

										__magic_name__						:						Any										=model(__snake_case     ,			token_type_ids=__snake_case     ,			labels=__snake_case			)
										self.parent.assertEqual(result.loss.shape     ,			()			)
										self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.vocab_size)			)

						def      A__    (		self :Optional[int]			):
										'''simple docstring'''

										__magic_name__						:						Tuple										=self.prepare_config_and_inputs()

										(
										    (
										    __magic_name__
										)					,			(
										    __magic_name__
										)					,			(
										    __magic_name__
										)					,			(
										    __magic_name__
										)					,			(
										    __magic_name__
										)					,			(
										    __magic_name__
										)					,			(
										    __magic_name__
										)					,			(
										    __magic_name__
										)					,			(
										    __magic_name__
										)					,			
										)						:						Any										=config_and_inputs

										__magic_name__						:						str										={"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """head_mask""": head_mask}

										return config, inputs_dict


						def      A__    (		self :Union[str, Any]     ,			__snake_case :Union[str, Any]     ,			__snake_case :str     ,			__snake_case :Optional[int]     ,			__snake_case :str     ,			*__snake_case :List[str]			):
										'''simple docstring'''

										__magic_name__						:						Tuple										=self.num_labels
										__magic_name__						:						str										=CTRLForSequenceClassification(__snake_case			)
										model.to(__snake_case			)
										model.eval()
										__magic_name__						:						str										=ids_tensor([self.batch_size]     ,			self.type_sequence_label_size			)
										__magic_name__						:						Any										=model(__snake_case     ,			token_type_ids=__snake_case     ,			labels=__snake_case			)
										self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.num_labels)			)




@require_torch
class 		__A   (      UpperCamelCase__       ,    UpperCamelCase__       ,    UpperCamelCase__       ,    unittest.TestCase      ):
						UpperCamelCase      =       (CTRLModel, CTRLLMHeadModel, CTRLForSequenceClassification) if is_torch_available() else ()
						UpperCamelCase      =       (CTRLLMHeadModel,) if is_torch_available() else ()
						UpperCamelCase      =       (
						    {
						        """feature-extraction""": CTRLModel,
						        """text-classification""": CTRLForSequenceClassification,
						        """text-generation""": CTRLLMHeadModel,
						        """zero-shot""": CTRLForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						UpperCamelCase      =       True
						UpperCamelCase      =       False
						UpperCamelCase      =       False

						def      A__    (		self :Any     ,			__snake_case :List[Any]     ,			__snake_case :List[str]     ,			__snake_case :List[Any]     ,			__snake_case :List[str]     ,			__snake_case :int			):
										'''simple docstring'''

										if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
														# Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
														# `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
														# config could not be created.
														return True

										return False

						def      A__    (		self :str			):
										'''simple docstring'''

										__magic_name__						:						Any										=CTRLModelTester(self			)
										__magic_name__						:						Optional[int]										=ConfigTester(self     ,			config_class=__snake_case     ,			n_embd=37			)

						def      A__    (		self :Any			):
										'''simple docstring'''

										super().tearDown()
										# clean-up as much as possible GPU memory occupied by PyTorch
										gc.collect()
										torch.cuda.empty_cache()

						def      A__    (		self :Any			):
										'''simple docstring'''

										self.config_tester.run_common_tests()

						def      A__    (		self :Union[str, Any]			):
										'''simple docstring'''

										__magic_name__						:						Tuple										=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_ctrl_model(*__snake_case			)

						def      A__    (		self :Dict			):
										'''simple docstring'''

										__magic_name__						:						Dict										=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_lm_head_model(*__snake_case			)

						@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""			)
						def      A__    (		self :Tuple			):
										'''simple docstring'''

										pass

						@slow
						def      A__    (		self :Any			):
										'''simple docstring'''

										for model_name in CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														__magic_name__						:						Tuple										=CTRLModel.from_pretrained(__snake_case			)
														self.assertIsNotNone(__snake_case			)


						@unittest.skip("""The model doesn't support left padding"""			)  # and it's not used enough to be worth fixing :)
						def      A__    (		self :List[Any]			):
										'''simple docstring'''

										pass




@require_torch
class 		__A   (      unittest.TestCase      ):

						def      A__    (		self :Dict			):
										'''simple docstring'''

										super().tearDown()
										# clean-up as much as possible GPU memory occupied by PyTorch
										gc.collect()
										torch.cuda.empty_cache()


						@slow
						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						int										=CTRLLMHeadModel.from_pretrained("""ctrl"""			)
										model.to(__snake_case			)
										__magic_name__						:						Optional[Any]										=torch.tensor(
										    [[1_18_59, 0, 16_11, 8]]     ,			dtype=torch.long     ,			device=__snake_case			)  # Legal the president is
										__magic_name__						:						Optional[int]										=[
										    1_18_59,
										    0,
										    16_11,
										    8,
										    5,
										    1_50,
										    2_64_49,
										    2,
										    19,
										    3_48,
										    4_69,
										    3,
										    25_95,
										    48,
										    2_07_40,
										    24_65_33,
										    24_65_33,
										    19,
										    30,
										    5,
										]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

										__magic_name__						:						str										=model.generate(__snake_case     ,			do_sample=__snake_case			)
										self.assertListEqual(output_ids[0].tolist()     ,			__snake_case			)










import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu




class 		__A   (      nn.Module      ):

						def __init__(		self :List[Any]			):
										'''simple docstring'''

										super().__init__()
										__magic_name__						:						Tuple										=nn.Linear(3     ,			4			)
										__magic_name__						:						Union[str, Any]										=nn.BatchNormad(4			)
										__magic_name__						:						List[str]										=nn.Linear(4     ,			5			)


						def      A__    (		self :Dict     ,			__snake_case :Tuple			):
										'''simple docstring'''

										return self.lineara(self.batchnorm(self.lineara(__snake_case			)			)			)




class 		__A   (      UpperCamelCase__      ):

						def      A__    (		self :Any     ,			__snake_case :Optional[Any]     ,			*__snake_case :List[Any]     ,			**__snake_case :Any			):
										'''simple docstring'''

										return (args[0] + 1,) + args[1:], kwargs




class 		__A   (      UpperCamelCase__      ):

						def      A__    (		self :List[str]     ,			__snake_case :Tuple     ,			__snake_case :Union[str, Any]			):
										'''simple docstring'''

										return output + 1




class 		__A   (      unittest.TestCase      ):

						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						int										=ModelForTest()
										__magic_name__						:						Tuple										=ModelHook()

										add_hook_to_module(__snake_case     ,			__snake_case			)
										self.assertEqual(test_model._hf_hook     ,			__snake_case			)
										self.assertTrue(hasattr(__snake_case     ,			"""_old_forward"""			)			)

										# Check adding the hook did not change the name or the signature
										self.assertEqual(test_model.forward.__name__     ,			"""forward"""			)
										self.assertListEqual(list(inspect.signature(test_model.forward			).parameters			)     ,			["""x"""]			)

										remove_hook_from_module(__snake_case			)
										self.assertFalse(hasattr(__snake_case     ,			"""_hf_hook"""			)			)
										self.assertFalse(hasattr(__snake_case     ,			"""_old_forward"""			)			)

						def      A__    (		self :Tuple			):
										'''simple docstring'''

										__magic_name__						:						int										=ModelForTest()
										__magic_name__						:						List[str]										=ModelHook()

										add_hook_to_module(__snake_case     ,			__snake_case			)
										add_hook_to_module(__snake_case     ,			__snake_case     ,			append=__snake_case			)

										self.assertEqual(isinstance(test_model._hf_hook     ,			__snake_case			)     ,			__snake_case			)
										self.assertEqual(len(test_model._hf_hook.hooks			)     ,			2			)
										self.assertTrue(hasattr(__snake_case     ,			"""_old_forward"""			)			)

										# Check adding the hook did not change the name or the signature
										self.assertEqual(test_model.forward.__name__     ,			"""forward"""			)
										self.assertListEqual(list(inspect.signature(test_model.forward			).parameters			)     ,			["""x"""]			)

										remove_hook_from_module(__snake_case			)
										self.assertFalse(hasattr(__snake_case     ,			"""_hf_hook"""			)			)
										self.assertFalse(hasattr(__snake_case     ,			"""_old_forward"""			)			)

						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						Any										=ModelForTest()
										__magic_name__						:						Any										=torch.randn(2     ,			3			)
										__magic_name__						:						Any										=test_model(x + 1			)
										__magic_name__						:						Optional[Any]										=test_model(x + 2			)

										__magic_name__						:						int										=PreForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						int										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			__snake_case     ,			atol=1E-5			)			)

										# Attaching a hook to a model when it already has one replaces, does not chain
										__magic_name__						:						str										=PreForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						List[str]										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			__snake_case     ,			atol=1E-5			)			)

										# You need to use the sequential hook to chain two or more hooks
										__magic_name__						:						Optional[Any]										=SequentialHook(PreForwardHook()     ,			PreForwardHook()			)
										add_hook_to_module(__snake_case     ,			__snake_case			)

										__magic_name__						:						Any										=test_model(__snake_case			)
										assert torch.allclose(__snake_case     ,			__snake_case     ,			atol=1E-5			)

						def      A__    (		self :Any			):
										'''simple docstring'''

										__magic_name__						:						Optional[Any]										=ModelForTest()
										__magic_name__						:						Dict										=torch.randn(2     ,			3			)
										__magic_name__						:						Any										=test_model(__snake_case			)

										__magic_name__						:						Dict										=PostForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						Any										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			output + 1     ,			atol=1E-5			)			)

										# Attaching a hook to a model when it already has one replaces, does not chain
										__magic_name__						:						Union[str, Any]										=PostForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						Optional[int]										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			output + 1     ,			atol=1E-5			)			)

										# You need to use the sequential hook to chain two or more hooks
										__magic_name__						:						Union[str, Any]										=SequentialHook(PostForwardHook()     ,			PostForwardHook()			)
										add_hook_to_module(__snake_case     ,			__snake_case			)

										__magic_name__						:						Union[str, Any]										=test_model(__snake_case			)
										assert torch.allclose(__snake_case     ,			output + 2     ,			atol=1E-5			)

						def      A__    (		self :Tuple			):
										'''simple docstring'''

										__magic_name__						:						Tuple										=ModelForTest()
										__magic_name__						:						int										=torch.randn(2     ,			3			)
										__magic_name__						:						Union[str, Any]										=test_model(__snake_case			)

										__magic_name__						:						Union[str, Any]										=PostForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						Dict										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			output + 1			)			)
										self.assertTrue(outputa.requires_grad			)

										__magic_name__						:						Any										=True
										__magic_name__						:						Any										=test_model(__snake_case			)
										self.assertFalse(outputa.requires_grad			)

						@require_multi_gpu
						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						Optional[Any]										=ModelForTest()
										# Everything is on CPU
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# This will move each submodule on different devices
										add_hook_to_module(model.lineara     ,			AlignDevicesHook(execution_device=0			)			)
										add_hook_to_module(model.batchnorm     ,			AlignDevicesHook(execution_device=0			)			)
										add_hook_to_module(model.lineara     ,			AlignDevicesHook(execution_device=1			)			)

										self.assertEqual(model.lineara.weight.device     ,			torch.device(0			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device(0			)			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			torch.device(0			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device(1			)			)

										# We can still make a forward pass. The input does not need to be on any particular device
										__magic_name__						:						List[Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						Optional[Any]										=model(__snake_case			)
										self.assertEqual(output.device     ,			torch.device(1			)			)

										# We can add a general hook to put back output on same device as input.
										add_hook_to_module(__snake_case     ,			AlignDevicesHook(io_same_device=__snake_case			)			)
										__magic_name__						:						int										=torch.randn(2     ,			3			).to(0			)
										__magic_name__						:						Optional[int]										=model(__snake_case			)
										self.assertEqual(output.device     ,			torch.device(0			)			)

						def      A__    (		self :List[Any]			):
										'''simple docstring'''

										__magic_name__						:						int										=ModelForTest()

										# Everything is on CPU
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# This will move each submodule on different devices
										__magic_name__						:						int										={"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

										add_hook_to_module(model.lineara     ,			AlignDevicesHook(**__snake_case			)			)
										add_hook_to_module(model.batchnorm     ,			AlignDevicesHook(**__snake_case			)			)
										add_hook_to_module(model.lineara     ,			AlignDevicesHook(**__snake_case			)			)

										# Parameters have been offloaded, so on the meta device
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										# Buffers are not included in the offload by default, so are on the execution device
										__magic_name__						:						Optional[int]										=torch.device(hook_kwargs["""execution_device"""]			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			__snake_case			)

										__magic_name__						:						Union[str, Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						Optional[int]										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_module(model.lineara			)
										remove_hook_from_module(model.batchnorm			)
										remove_hook_from_module(model.lineara			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# Now test with buffers included in the offload
										__magic_name__						:						Tuple										={
										    """execution_device""": 0 if torch.cuda.is_available() else """cpu""",
										    """offload""": True,
										    """offload_buffers""": True,
										}

										add_hook_to_module(model.lineara     ,			AlignDevicesHook(**__snake_case			)			)
										add_hook_to_module(model.batchnorm     ,			AlignDevicesHook(**__snake_case			)			)
										add_hook_to_module(model.lineara     ,			AlignDevicesHook(**__snake_case			)			)

										# Parameters have been offloaded, so on the meta device, buffers included
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			torch.device("""meta"""			)			)

										__magic_name__						:						Tuple										=torch.randn(2     ,			3			)
										__magic_name__						:						int										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_module(model.lineara			)
										remove_hook_from_module(model.batchnorm			)
										remove_hook_from_module(model.lineara			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

						def      A__    (		self :List[Any]			):
										'''simple docstring'''

										__magic_name__						:						Any										=ModelForTest()

										# Everything is on CPU
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# This will move each submodule on different devices
										__magic_name__						:						str										=0 if torch.cuda.is_available() else """cpu"""
										attach_align_device_hook(__snake_case     ,			execution_device=__snake_case     ,			offload=__snake_case			)

										# Parameters have been offloaded, so on the meta device
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										# Buffers are not included in the offload by default, so are on the execution device
										__magic_name__						:						Optional[Any]										=torch.device(__snake_case			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			__snake_case			)

										__magic_name__						:						List[Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						str										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(__snake_case			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# Now test with buffers included in the offload
										attach_align_device_hook(__snake_case     ,			execution_device=__snake_case     ,			offload=__snake_case     ,			offload_buffers=__snake_case			)

										# Parameters have been offloaded, so on the meta device, buffers included
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			torch.device("""meta"""			)			)

										__magic_name__						:						Optional[int]										=torch.randn(2     ,			3			)
										__magic_name__						:						Union[str, Any]										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(__snake_case			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)


						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						Dict										=ModelForTest()

										# Everything is on CPU
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# This will move each submodule on different devices
										__magic_name__						:						List[str]										=0 if torch.cuda.is_available() else """cpu"""
										attach_align_device_hook(
										    __snake_case     ,			execution_device=__snake_case     ,			offload=__snake_case     ,			weights_map=model.state_dict()			)

										# Parameters have been offloaded, so on the meta device
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										# Buffers are not included in the offload by default, so are on the execution device
										__magic_name__						:						Optional[Any]										=torch.device(__snake_case			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			__snake_case			)

										__magic_name__						:						int										=torch.randn(2     ,			3			)
										__magic_name__						:						Any										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(__snake_case			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# Now test with buffers included in the offload
										attach_align_device_hook(
										    __snake_case     ,			execution_device=__snake_case     ,			offload=__snake_case     ,			weights_map=model.state_dict()     ,			offload_buffers=__snake_case     ,			)

										# Parameters have been offloaded, so on the meta device, buffers included
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			torch.device("""meta"""			)			)

										__magic_name__						:						List[Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						str										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(__snake_case			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)




