


import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import (
    BitConfig,
    ViTHybridConfig,
    ViTHybridForImageClassification,
    ViTHybridImageProcessor,
    ViTHybridModel,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase           =		logging.get_logger(__name__)







def        _lowerCamelCase(	lowercase__   , lowercase__=False	)						->	Any:

		'''simple docstring'''







		__lowercase=			[]

		# fmt: off
		# stem:
		rename_keys.append(('cls_token', 'vit.embeddings.cls_token')	)
		rename_keys.append(('pos_embed', 'vit.embeddings.position_embeddings')	)

		rename_keys.append(('patch_embed.proj.weight', 'vit.embeddings.patch_embeddings.projection.weight')	)
		rename_keys.append(('patch_embed.proj.bias', 'vit.embeddings.patch_embeddings.projection.bias')	)

		# backbone
		rename_keys.append(('patch_embed.backbone.stem.conv.weight', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.convolution.weight')	)
		rename_keys.append(('patch_embed.backbone.stem.norm.weight', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.weight')	)
		rename_keys.append(('patch_embed.backbone.stem.norm.bias', 'vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.bias')	)

		for stage_idx in range(len(config.backbone_config.depths	)	):
				for layer_idx in range(config.backbone_config.depths[stage_idx]	):
						rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv1.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv1.weight')	)
						rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.weight')	)
						rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.bias')	)
						rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv2.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv2.weight')	)
						rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.weight')	)
						rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.bias')	)
						rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv3.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv3.weight')	)
						rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.weight')	)
						rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.bias')	)

				rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.conv.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.conv.weight')	)
				rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.weight')	)
				rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.bias')	)

		# transformer encoder
		for i in range(config.num_hidden_layers	):
				# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
				rename_keys.append((F'blocks.{i}.norm1.weight', F'vit.encoder.layer.{i}.layernorm_before.weight')	)
				rename_keys.append((F'blocks.{i}.norm1.bias', F'vit.encoder.layer.{i}.layernorm_before.bias')	)
				rename_keys.append((F'blocks.{i}.attn.proj.weight', F'vit.encoder.layer.{i}.attention.output.dense.weight')	)
				rename_keys.append((F'blocks.{i}.attn.proj.bias', F'vit.encoder.layer.{i}.attention.output.dense.bias')	)
				rename_keys.append((F'blocks.{i}.norm2.weight', F'vit.encoder.layer.{i}.layernorm_after.weight')	)
				rename_keys.append((F'blocks.{i}.norm2.bias', F'vit.encoder.layer.{i}.layernorm_after.bias')	)
				rename_keys.append((F'blocks.{i}.mlp.fc1.weight', F'vit.encoder.layer.{i}.intermediate.dense.weight')	)
				rename_keys.append((F'blocks.{i}.mlp.fc1.bias', F'vit.encoder.layer.{i}.intermediate.dense.bias')	)
				rename_keys.append((F'blocks.{i}.mlp.fc2.weight', F'vit.encoder.layer.{i}.output.dense.weight')	)
				rename_keys.append((F'blocks.{i}.mlp.fc2.bias', F'vit.encoder.layer.{i}.output.dense.bias')	)

		if base_model:
				# layernorm + pooler
				rename_keys.extend(
				    [
				        ('norm.weight', 'layernorm.weight'),
				        ('norm.bias', 'layernorm.bias'),
				        ('pre_logits.fc.weight', 'pooler.dense.weight'),
				        ('pre_logits.fc.bias', 'pooler.dense.bias'),
				    ]	)

				# if just the base model, we should remove "vit" from all keys that start with "vit"
				__lowercase=			[(pair[0], pair[1][4:]) if pair[1].startswith('vit'	) else pair for pair in rename_keys]
		else:
				# layernorm + classification head
				rename_keys.extend(
				    [
				        ('norm.weight', 'vit.layernorm.weight'),
				        ('norm.bias', 'vit.layernorm.bias'),
				        ('head.weight', 'classifier.weight'),
				        ('head.bias', 'classifier.bias'),
				    ]	)
		# fmt: on

		return rename_keys







def        _lowerCamelCase(	lowercase__   , lowercase__   , lowercase__=False	)						->	Union[str, Any]:

		'''simple docstring'''







		for i in range(config.num_hidden_layers	):
				if base_model:
						__lowercase=			''
				else:
						__lowercase=			'vit.'
				# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
				__lowercase=			state_dict.pop(F'blocks.{i}.attn.qkv.weight'	)
				__lowercase=			state_dict.pop(F'blocks.{i}.attn.qkv.bias'	)
				# next, add query, keys and values (in that order) to the state dict
				__lowercase=			in_proj_weight[
				    : config.hidden_size, :
				]
				__lowercase=			in_proj_bias[: config.hidden_size]
				__lowercase=			in_proj_weight[
				    config.hidden_size : config.hidden_size * 2, :
				]
				__lowercase=			in_proj_bias[
				    config.hidden_size : config.hidden_size * 2
				]
				__lowercase=			in_proj_weight[
				    -config.hidden_size :, :
				]
				__lowercase=			in_proj_bias[-config.hidden_size :]







def        _lowerCamelCase(	lowercase__	)						->	str:

		'''simple docstring'''







		__lowercase=			['head.weight', 'head.bias']
		for k in ignore_keys:
				state_dict.pop(lowercase__   , lowercase__	)







def        _lowerCamelCase(	lowercase__   , lowercase__   , lowercase__	)						->	Dict:

		'''simple docstring'''







		__lowercase=			dct.pop(lowercase__	)
		__lowercase=			val







def        _lowerCamelCase(	)						->	str:

		'''simple docstring'''







		__lowercase=			'http://images.cocodataset.org/val2017/000000039769.jpg'
		__lowercase=			Image.open(requests.get(lowercase__   , stream=lowercase__	).raw	)
		return im







@torch.no_grad()
def        _lowerCamelCase(	lowercase__   , lowercase__   , lowercase__=False	)						->	List[str]:

		'''simple docstring'''







		__lowercase=			BitConfig(
		    global_padding='same'   , layer_type='bottleneck'   , depths=(3, 4, 9)   , out_features=['stage3']   , embedding_dynamic_padding=lowercase__   , )
		__lowercase=			ViTHybridConfig(backbone_config=lowercase__   , image_size=3_8_4   , num_labels=1_0_0_0	)
		__lowercase=			False

		# load original model from timm
		__lowercase=			timm.create_model(lowercase__   , pretrained=lowercase__	)
		timm_model.eval()

		# load state_dict of original model, remove and rename some keys
		__lowercase=			timm_model.state_dict()
		if base_model:
				remove_classification_head_(lowercase__	)
		__lowercase=			create_rename_keys(lowercase__   , lowercase__	)
		for src, dest in rename_keys:
				rename_key(lowercase__   , lowercase__   , lowercase__	)
		read_in_q_k_v(lowercase__   , lowercase__   , lowercase__	)

		__lowercase=			'huggingface/label-files'
		__lowercase=			'imagenet-1k-id2label.json'
		__lowercase=			json.load(open(hf_hub_download(lowercase__   , lowercase__   , repo_type='dataset'	)   , 'r'	)	)
		__lowercase=			{int(lowercase__	): v for k, v in idalabel.items()}
		__lowercase=			idalabel
		__lowercase=			{v: k for k, v in idalabel.items()}

		# load HuggingFace model
		if vit_name[-5:] == "in21k":
				__lowercase=			ViTHybridModel(lowercase__	).eval()
		else:
				__lowercase=			ViTHybridForImageClassification(lowercase__	).eval()
		model.load_state_dict(lowercase__	)

		# create image processor
		__lowercase=			create_transform(**resolve_data_config({}   , model=lowercase__	)	)
		__lowercase=			transform.transforms

		__lowercase=			{
		    'bilinear': PILImageResampling.BILINEAR,
		    'bicubic': PILImageResampling.BICUBIC,
		    'nearest': PILImageResampling.NEAREST,
		}

		__lowercase=			ViTHybridImageProcessor(
		    do_resize=lowercase__   , size={'shortest_edge': timm_transforms[0].size}   , resample=pillow_resamplings[timm_transforms[0].interpolation.value]   , do_center_crop=lowercase__   , crop_size={'height': timm_transforms[1].size[0], 'width': timm_transforms[1].size[1]}   , do_normalize=lowercase__   , image_mean=timm_transforms[-1].mean.tolist()   , image_std=timm_transforms[-1].std.tolist()   , )

		__lowercase=			prepare_img()
		__lowercase=			transform(lowercase__	).unsqueeze(0	)
		__lowercase=			processor(lowercase__   , return_tensors='pt'	).pixel_values

		# verify pixel values
		assert torch.allclose(lowercase__   , lowercase__	)

		# verify logits
		with torch.no_grad():
				__lowercase=			model(lowercase__	)
				__lowercase=			outputs.logits

		print('Predicted class:'   , logits.argmax(-1	).item()	)
		if base_model:
				__lowercase=			timm_model.forward_features(lowercase__	)
				assert timm_pooled_output.shape == outputs.pooler_output.shape
				assert torch.allclose(lowercase__   , outputs.pooler_output   , atol=1E-3	)
		else:
				__lowercase=			timm_model(lowercase__	)
				assert timm_logits.shape == outputs.logits.shape
				assert torch.allclose(lowercase__   , outputs.logits   , atol=1E-3	)
		print('Looks ok!'	)

		if pytorch_dump_folder_path is not None:
				Path(lowercase__	).mkdir(exist_ok=lowercase__	)
				print(F'Saving model {vit_name} to {pytorch_dump_folder_path}'	)
				model.save_pretrained(lowercase__	)
				print(F'Saving processor to {pytorch_dump_folder_path}'	)
				processor.save_pretrained(lowercase__	)

		if push_to_hub:
				print(F'Pushing model and processor to the hub {vit_name}'	)
				model.push_to_hub(F'ybelkada/{vit_name}'	)
				processor.push_to_hub(F'ybelkada/{vit_name}'	)


if __name__ == "__main__":
	lowerCAmelCase           =		argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '''--vit_name''',
	    default='''vit_base_r50_s16_384''',
	    type=str,
	    help='''Name of the hybrid ViT timm model you\'d like to convert.''',
	)
	parser.add_argument(
	    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
	)
	parser.add_argument(
	    '''--push_to_hub''', action='''store_true''', help='''Whether to upload the model to the HuggingFace hub.'''
	)

	lowerCAmelCase           =		parser.parse_args()
	convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)










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					):
									super().__init__()
									__lowercase=			nn.Linear(3      ,				4					)
									__lowercase=			nn.BatchNormad(4					)
									__lowercase=			nn.Linear(4      ,				5					)







							def 			_A       (self      ,				lowerCAmelCase					):
									return self.lineara(self.batchnorm(self.lineara(lowerCAmelCase					)					)					)





class       A    (   A_ ):


							def 			_A       (self      ,				lowerCAmelCase      ,				*lowerCAmelCase      ,				**lowerCAmelCase					):
									return (args[0] + 1,) + args[1:], kwargs





class       A    (   A_ ):


							def 			_A       (self      ,				lowerCAmelCase      ,				lowerCAmelCase					):
									return output + 1





class       A    (   unittest.TestCase ):


							def 			_A       (self					):
									__lowercase=			ModelForTest()
									__lowercase=			ModelHook()

									add_hook_to_module(lowerCAmelCase      ,				lowerCAmelCase					)
									self.assertEqual(test_model._hf_hook      ,				lowerCAmelCase					)
									self.assertTrue(hasattr(lowerCAmelCase      ,				'_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(lowerCAmelCase					)
									self.assertFalse(hasattr(lowerCAmelCase      ,				'_hf_hook'					)					)
									self.assertFalse(hasattr(lowerCAmelCase      ,				'_old_forward'					)					)


							def 			_A       (self					):
									__lowercase=			ModelForTest()
									__lowercase=			ModelHook()

									add_hook_to_module(lowerCAmelCase      ,				lowerCAmelCase					)
									add_hook_to_module(lowerCAmelCase      ,				lowerCAmelCase      ,				append=lowerCAmelCase					)

									self.assertEqual(isinstance(test_model._hf_hook      ,				lowerCAmelCase					)      ,				lowerCAmelCase					)
									self.assertEqual(len(test_model._hf_hook.hooks					)      ,				2					)
									self.assertTrue(hasattr(lowerCAmelCase      ,				'_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(lowerCAmelCase					)
									self.assertFalse(hasattr(lowerCAmelCase      ,				'_hf_hook'					)					)
									self.assertFalse(hasattr(lowerCAmelCase      ,				'_old_forward'					)					)


							def 			_A       (self					):
									__lowercase=			ModelForTest()
									__lowercase=			torch.randn(2      ,				3					)
									__lowercase=			test_model(x + 1					)
									__lowercase=			test_model(x + 2					)

									__lowercase=			PreForwardHook()
									add_hook_to_module(lowerCAmelCase      ,				lowerCAmelCase					)
									__lowercase=			test_model(lowerCAmelCase					)
									self.assertTrue(torch.allclose(lowerCAmelCase      ,				lowerCAmelCase      ,				atol=1E-5					)					)

									# Attaching a hook to a model when it already has one replaces, does not chain
									__lowercase=			PreForwardHook()
									add_hook_to_module(lowerCAmelCase      ,				lowerCAmelCase					)
									__lowercase=			test_model(lowerCAmelCase					)
									self.assertTrue(torch.allclose(lowerCAmelCase      ,				lowerCAmelCase      ,				atol=1E-5					)					)

									# You need to use the sequential hook to chain two or more hooks
									__lowercase=			SequentialHook(PreForwardHook()      ,				PreForwardHook()					)
									add_hook_to_module(lowerCAmelCase      ,				lowerCAmelCase					)

									__lowercase=			test_model(lowerCAmelCase					)
									assert torch.allclose(lowerCAmelCase      ,				lowerCAmelCase      ,				atol=1E-5					)


							def 			_A       (self					):
									__lowercase=			ModelForTest()
									__lowercase=			torch.randn(2      ,				3					)
									__lowercase=			test_model(lowerCAmelCase					)

									__lowercase=			PostForwardHook()
									add_hook_to_module(lowerCAmelCase      ,				lowerCAmelCase					)
									__lowercase=			test_model(lowerCAmelCase					)
									self.assertTrue(torch.allclose(lowerCAmelCase      ,				output + 1      ,				atol=1E-5					)					)

									# Attaching a hook to a model when it already has one replaces, does not chain
									__lowercase=			PostForwardHook()
									add_hook_to_module(lowerCAmelCase      ,				lowerCAmelCase					)
									__lowercase=			test_model(lowerCAmelCase					)
									self.assertTrue(torch.allclose(lowerCAmelCase      ,				output + 1      ,				atol=1E-5					)					)

									# You need to use the sequential hook to chain two or more hooks
									__lowercase=			SequentialHook(PostForwardHook()      ,				PostForwardHook()					)
									add_hook_to_module(lowerCAmelCase      ,				lowerCAmelCase					)

									__lowercase=			test_model(lowerCAmelCase					)
									assert torch.allclose(lowerCAmelCase      ,				output + 2      ,				atol=1E-5					)


							def 			_A       (self					):
									__lowercase=			ModelForTest()
									__lowercase=			torch.randn(2      ,				3					)
									__lowercase=			test_model(lowerCAmelCase					)

									__lowercase=			PostForwardHook()
									add_hook_to_module(lowerCAmelCase      ,				lowerCAmelCase					)
									__lowercase=			test_model(lowerCAmelCase					)
									self.assertTrue(torch.allclose(lowerCAmelCase      ,				output + 1					)					)
									self.assertTrue(outputa.requires_grad					)

									__lowercase=			True
									__lowercase=			test_model(lowerCAmelCase					)
									self.assertFalse(outputa.requires_grad					)


							@require_multi_gpu
							def 			_A       (self					):
									__lowercase=			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
									__lowercase=			torch.randn(2      ,				3					)
									__lowercase=			model(lowerCAmelCase					)
									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(lowerCAmelCase      ,				AlignDevicesHook(io_same_device=lowerCAmelCase					)					)
									__lowercase=			torch.randn(2      ,				3					).to(0					)
									__lowercase=			model(lowerCAmelCase					)
									self.assertEqual(output.device      ,				torch.device(0					)					)


							def 			_A       (self					):
									__lowercase=			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
									__lowercase=			{'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

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

									# 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
									__lowercase=			torch.device(hook_kwargs['execution_device']					)
									self.assertEqual(model.batchnorm.running_mean.device      ,				lowerCAmelCase					)

									__lowercase=			torch.randn(2      ,				3					)
									__lowercase=			model(lowerCAmelCase					)
									self.assertEqual(output.device      ,				lowerCAmelCase					)

									# 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
									__lowercase=			{
									    'execution_device': 0 if torch.cuda.is_available() else 'cpu',
									    'offload': True,
									    'offload_buffers': True,
									}

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

									# 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'					)					)

									__lowercase=			torch.randn(2      ,				3					)
									__lowercase=			model(lowerCAmelCase					)
									self.assertEqual(output.device      ,				lowerCAmelCase					)

									# 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					):
									__lowercase=			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
									__lowercase=			0 if torch.cuda.is_available() else 'cpu'
									attach_align_device_hook(lowerCAmelCase      ,				execution_device=lowerCAmelCase      ,				offload=lowerCAmelCase					)

									# 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
									__lowercase=			torch.device(lowerCAmelCase					)
									self.assertEqual(model.batchnorm.running_mean.device      ,				lowerCAmelCase					)

									__lowercase=			torch.randn(2      ,				3					)
									__lowercase=			model(lowerCAmelCase					)
									self.assertEqual(output.device      ,				lowerCAmelCase					)

									# Removing hooks loads back the weights in the model.
									remove_hook_from_submodules(lowerCAmelCase					)
									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(lowerCAmelCase      ,				execution_device=lowerCAmelCase      ,				offload=lowerCAmelCase      ,				offload_buffers=lowerCAmelCase					)

									# 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'					)					)

									__lowercase=			torch.randn(2      ,				3					)
									__lowercase=			model(lowerCAmelCase					)
									self.assertEqual(output.device      ,				lowerCAmelCase					)

									# Removing hooks loads back the weights in the model.
									remove_hook_from_submodules(lowerCAmelCase					)
									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					):
									__lowercase=			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
									__lowercase=			0 if torch.cuda.is_available() else 'cpu'
									attach_align_device_hook(
									    lowerCAmelCase      ,				execution_device=lowerCAmelCase      ,				offload=lowerCAmelCase      ,				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
									__lowercase=			torch.device(lowerCAmelCase					)
									self.assertEqual(model.batchnorm.running_mean.device      ,				lowerCAmelCase					)

									__lowercase=			torch.randn(2      ,				3					)
									__lowercase=			model(lowerCAmelCase					)
									self.assertEqual(output.device      ,				lowerCAmelCase					)

									# Removing hooks loads back the weights in the model.
									remove_hook_from_submodules(lowerCAmelCase					)
									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(
									    lowerCAmelCase      ,				execution_device=lowerCAmelCase      ,				offload=lowerCAmelCase      ,				weights_map=model.state_dict()      ,				offload_buffers=lowerCAmelCase      ,				)

									# 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'					)					)

									__lowercase=			torch.randn(2      ,				3					)
									__lowercase=			model(lowerCAmelCase					)
									self.assertEqual(output.device      ,				lowerCAmelCase					)

									# Removing hooks loads back the weights in the model.
									remove_hook_from_submodules(lowerCAmelCase					)
									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'					)					)






