




import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

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


if is_torch_available():
							import torch

							from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM




@require_torch
class 		__A   :

						def __init__(		self :Optional[Any]     ,			__snake_case :List[str]     ,			__snake_case :Any=99     ,			__snake_case :Optional[int]=13     ,			__snake_case :Any=16     ,			__snake_case :Any=7     ,			__snake_case :Tuple=True     ,			__snake_case :List[Any]=True     ,			__snake_case :Union[str, Any]=True     ,			__snake_case :Dict=False     ,			__snake_case :str=True     ,			__snake_case :Optional[Any]=2     ,			__snake_case :Optional[int]=32     ,			__snake_case :int=4     ,			__snake_case :str=4     ,			__snake_case :Optional[int]=30     ,			__snake_case :str=0     ,			__snake_case :Optional[Any]=1     ,			__snake_case :Optional[int]=2     ,			__snake_case :List[str]=None     ,			):
										'''simple docstring'''

										__magic_name__						:						Union[str, Any]										=parent
										__magic_name__						:						str										=batch_size
										__magic_name__						:						List[str]										=decoder_seq_length
										# For common tests
										__magic_name__						:						Optional[Any]										=self.decoder_seq_length
										__magic_name__						:						Any										=is_training
										__magic_name__						:						Any										=use_attention_mask
										__magic_name__						:						List[str]										=use_labels

										__magic_name__						:						List[str]										=vocab_size
										__magic_name__						:						str										=d_model
										__magic_name__						:						Optional[int]										=d_model
										__magic_name__						:						str										=decoder_layers
										__magic_name__						:						Optional[Any]										=decoder_layers
										__magic_name__						:						Dict										=decoder_ffn_dim
										__magic_name__						:						Tuple										=decoder_attention_heads
										__magic_name__						:						Optional[Any]										=decoder_attention_heads
										__magic_name__						:						Union[str, Any]										=eos_token_id
										__magic_name__						:						Optional[Any]										=bos_token_id
										__magic_name__						:						Optional[Any]										=pad_token_id
										__magic_name__						:						Any										=decoder_start_token_id
										__magic_name__						:						str										=use_cache
										__magic_name__						:						Optional[int]										=max_position_embeddings

										__magic_name__						:						Any										=None
										__magic_name__						:						Dict										=decoder_seq_length
										__magic_name__						:						Union[str, Any]										=2
										__magic_name__						:						List[str]										=1

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

										__magic_name__						:						Union[str, Any]										=ids_tensor([self.batch_size, self.decoder_seq_length]     ,			self.vocab_size			)

										__magic_name__						:						Any										=None
										if self.use_attention_mask:
														__magic_name__						:						int										=ids_tensor([self.batch_size, self.decoder_seq_length]     ,			vocab_size=2			)

										__magic_name__						:						Optional[int]										=None
										if self.use_labels:
														__magic_name__						:						List[str]										=ids_tensor([self.batch_size, self.decoder_seq_length]     ,			self.vocab_size			)

										__magic_name__						:						Tuple										=TrOCRConfig(
										    vocab_size=self.vocab_size     ,			d_model=self.d_model     ,			decoder_layers=self.decoder_layers     ,			decoder_ffn_dim=self.decoder_ffn_dim     ,			decoder_attention_heads=self.decoder_attention_heads     ,			eos_token_id=self.eos_token_id     ,			bos_token_id=self.bos_token_id     ,			use_cache=self.use_cache     ,			pad_token_id=self.pad_token_id     ,			decoder_start_token_id=self.decoder_start_token_id     ,			max_position_embeddings=self.max_position_embeddings     ,			)

										return (config, input_ids, attention_mask, lm_labels)

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

										__magic_name__						:						Dict										=True
										__magic_name__						:						List[Any]										=TrOCRDecoder(config=__snake_case			).to(__snake_case			).eval()
										__magic_name__						:						List[Any]										=input_ids[:2]

										input_ids[input_ids == 0] += 1
										# first forward pass
										__magic_name__						:						str										=model(__snake_case     ,			use_cache=__snake_case			)
										__magic_name__						:						Optional[Any]										=model(__snake_case			)
										__magic_name__						:						str										=model(__snake_case     ,			use_cache=__snake_case			)

										self.parent.assertTrue(len(__snake_case			) == len(__snake_case			)			)
										self.parent.assertTrue(len(__snake_case			) == len(__snake_case			) + 1			)

										__magic_name__						:						str										=outputs["""past_key_values"""]

										# create hypothetical next token and extent to next_input_ids
										__magic_name__						:						Dict										=ids_tensor((2, 1)     ,			config.vocab_size - 1			) + 1

										# append to next input_ids and
										__magic_name__						:						List[Any]										=torch.cat([input_ids, next_tokens]     ,			dim=-1			)

										__magic_name__						:						Dict										=model(__snake_case			)["""last_hidden_state"""]
										__magic_name__						:						Union[str, Any]										=model(__snake_case     ,			past_key_values=__snake_case			)["""last_hidden_state"""]

										# select random slice
										__magic_name__						:						Dict										=ids_tensor((1,)     ,			output_from_past.shape[-1]			).item()
										__magic_name__						:						List[str]										=output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
										__magic_name__						:						List[str]										=output_from_past[:, 0, random_slice_idx].detach()

										# test that outputs are equal for slice
										assert torch.allclose(__snake_case     ,			__snake_case     ,			atol=1E-3			)


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

										__magic_name__						:						Tuple										=self.prepare_config_and_inputs()
										__magic_name__					,			__magic_name__					,			__magic_name__					,			__magic_name__						:						int										=config_and_inputs

										__magic_name__						:						Optional[int]										={"""input_ids""": input_ids, """attention_mask""": attention_mask}
										return config, inputs_dict




@require_torch
class 		__A   (      UpperCamelCase__       ,    UpperCamelCase__       ,    UpperCamelCase__       ,    unittest.TestCase      ):
						UpperCamelCase      =       (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
						UpperCamelCase      =       (TrOCRForCausalLM,) if is_torch_available() else ()
						UpperCamelCase      =       {"""text-generation""": TrOCRForCausalLM} if is_torch_available() else {}
						UpperCamelCase      =       True
						UpperCamelCase      =       False

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

										__magic_name__						:						Dict										=TrOCRStandaloneDecoderModelTester(self     ,			is_training=__snake_case			)
										__magic_name__						:						Optional[int]										=ConfigTester(self     ,			config_class=__snake_case			)

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

										pass

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

										pass

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

										pass

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

										self.config_tester.run_common_tests()

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

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

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

										return


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

										pass










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"""			)			)




