import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
from PIL import Image

from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase_				:		Optional[Any]       						=       logging.get_logger(__name__)
def 							lowerCAmelCase(	):
					__a							=   DetaConfig(
					    num_queries=900      ,				encoder_ffn_dim=2048      ,				decoder_ffn_dim=2048      ,				num_feature_levels=5      ,				assign_first_stage=snake_case__      ,				with_box_refine=snake_case__      ,				two_stage=snake_case__      ,				)

					# set labels
					__a							=   91
					__a							=   'huggingface/label-files'
					__a							=   'coco-detection-id2label.json'
					__a							=   json.load(open(cached_download(hf_hub_url(snake_case__      ,				snake_case__      ,				repo_type='dataset'						)						)      ,				'r'						)						)
					__a							=   {int(snake_case__						): v for k, v in idalabel.items()}
					__a							=   idalabel
					__a							=   {v: k for k, v in idalabel.items()}

					return config
def 							lowerCAmelCase(	__lowerCamelCase						):
					__a							=   []

					# stem
					# fmt: off
					rename_keys.append(('backbone.0.body.conv1.weight', 'model.backbone.model.embedder.embedder.convolution.weight')						)
					rename_keys.append(('backbone.0.body.bn1.weight', 'model.backbone.model.embedder.embedder.normalization.weight')						)
					rename_keys.append(('backbone.0.body.bn1.bias', 'model.backbone.model.embedder.embedder.normalization.bias')						)
					rename_keys.append(('backbone.0.body.bn1.running_mean', 'model.backbone.model.embedder.embedder.normalization.running_mean')						)
					rename_keys.append(('backbone.0.body.bn1.running_var', 'model.backbone.model.embedder.embedder.normalization.running_var')						)
					# stages
					for stage_idx in range(len(config.backbone_config.depths						)						):
										for layer_idx in range(config.backbone_config.depths[stage_idx]						):
															# shortcut
															if layer_idx == 0:
																				rename_keys.append(
																				    (
																				        f'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight''',
																				        f'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight''',
																				    )						)
																				rename_keys.append(
																				    (
																				        f'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight''',
																				        f'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight''',
																				    )						)
																				rename_keys.append(
																				    (
																				        f'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias''',
																				        f'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias''',
																				    )						)
																				rename_keys.append(
																				    (
																				        f'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean''',
																				        f'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean''',
																				    )						)
																				rename_keys.append(
																				    (
																				        f'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var''',
																				        f'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var''',
																				    )						)
															# 3 convs
															for i in range(3						):
																				rename_keys.append(
																				    (
																				        f'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight''',
																				        f'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight''',
																				    )						)
																				rename_keys.append(
																				    (
																				        f'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight''',
																				        f'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight''',
																				    )						)
																				rename_keys.append(
																				    (
																				        f'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias''',
																				        f'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias''',
																				    )						)
																				rename_keys.append(
																				    (
																				        f'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean''',
																				        f'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean''',
																				    )						)
																				rename_keys.append(
																				    (
																				        f'''backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var''',
																				        f'''model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var''',
																				    )						)
																				# transformer encoder
					for i in range(config.encoder_layers						):
										rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight''', f'''model.encoder.layers.{i}.self_attn.sampling_offsets.weight''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias''', f'''model.encoder.layers.{i}.self_attn.sampling_offsets.bias''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.attention_weights.weight''', f'''model.encoder.layers.{i}.self_attn.attention_weights.weight''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.attention_weights.bias''', f'''model.encoder.layers.{i}.self_attn.attention_weights.bias''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.value_proj.weight''', f'''model.encoder.layers.{i}.self_attn.value_proj.weight''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.value_proj.bias''', f'''model.encoder.layers.{i}.self_attn.value_proj.bias''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.output_proj.weight''', f'''model.encoder.layers.{i}.self_attn.output_proj.weight''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.self_attn.output_proj.bias''', f'''model.encoder.layers.{i}.self_attn.output_proj.bias''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.norm1.weight''', f'''model.encoder.layers.{i}.self_attn_layer_norm.weight''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.norm1.bias''', f'''model.encoder.layers.{i}.self_attn_layer_norm.bias''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.linear1.weight''', f'''model.encoder.layers.{i}.fc1.weight''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.linear1.bias''', f'''model.encoder.layers.{i}.fc1.bias''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.linear2.weight''', f'''model.encoder.layers.{i}.fc2.weight''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.linear2.bias''', f'''model.encoder.layers.{i}.fc2.bias''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.norm2.weight''', f'''model.encoder.layers.{i}.final_layer_norm.weight''')						)
										rename_keys.append((f'''transformer.encoder.layers.{i}.norm2.bias''', f'''model.encoder.layers.{i}.final_layer_norm.bias''')						)

					# transformer decoder
					for i in range(config.decoder_layers						):
										rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight''', f'''model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias''', f'''model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.attention_weights.weight''', f'''model.decoder.layers.{i}.encoder_attn.attention_weights.weight''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.attention_weights.bias''', f'''model.decoder.layers.{i}.encoder_attn.attention_weights.bias''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.value_proj.weight''', f'''model.decoder.layers.{i}.encoder_attn.value_proj.weight''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.value_proj.bias''', f'''model.decoder.layers.{i}.encoder_attn.value_proj.bias''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.output_proj.weight''', f'''model.decoder.layers.{i}.encoder_attn.output_proj.weight''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.cross_attn.output_proj.bias''', f'''model.decoder.layers.{i}.encoder_attn.output_proj.bias''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.norm1.weight''', f'''model.decoder.layers.{i}.encoder_attn_layer_norm.weight''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.norm1.bias''', f'''model.decoder.layers.{i}.encoder_attn_layer_norm.bias''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.self_attn.out_proj.weight''', f'''model.decoder.layers.{i}.self_attn.out_proj.weight''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.self_attn.out_proj.bias''', f'''model.decoder.layers.{i}.self_attn.out_proj.bias''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.norm2.weight''', f'''model.decoder.layers.{i}.self_attn_layer_norm.weight''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.norm2.bias''', f'''model.decoder.layers.{i}.self_attn_layer_norm.bias''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.linear1.weight''', f'''model.decoder.layers.{i}.fc1.weight''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.linear1.bias''', f'''model.decoder.layers.{i}.fc1.bias''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.linear2.weight''', f'''model.decoder.layers.{i}.fc2.weight''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.linear2.bias''', f'''model.decoder.layers.{i}.fc2.bias''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.norm3.weight''', f'''model.decoder.layers.{i}.final_layer_norm.weight''')						)
										rename_keys.append((f'''transformer.decoder.layers.{i}.norm3.bias''', f'''model.decoder.layers.{i}.final_layer_norm.bias''')						)

					# fmt: on

					return rename_keys
def 							lowerCAmelCase(	__lowerCamelCase      ,				__lowerCamelCase      ,				__lowerCamelCase						):
					__a							=   dct.pop(snake_case__						)
					__a							=   val
def 							lowerCAmelCase(	__lowerCamelCase      ,				__lowerCamelCase						):
					# transformer decoder self-attention layers
					__a							=   config.d_model
					for i in range(config.decoder_layers						):
										# read in weights + bias of input projection layer of self-attention
										__a							=   state_dict.pop(f'''transformer.decoder.layers.{i}.self_attn.in_proj_weight'''						)
										__a							=   state_dict.pop(f'''transformer.decoder.layers.{i}.self_attn.in_proj_bias'''						)
										# next, add query, keys and values (in that order) to the state dict
										__a							=   in_proj_weight[:hidden_size, :]
										__a							=   in_proj_bias[:hidden_size]
										__a							=   in_proj_weight[
										    hidden_size : hidden_size * 2, :
										]
										__a							=   in_proj_bias[hidden_size : hidden_size * 2]
										__a							=   in_proj_weight[-hidden_size:, :]
										__a							=   in_proj_bias[-hidden_size:]
def 							lowerCAmelCase(	):
					__a							=   'http://images.cocodataset.org/val2017/000000039769.jpg'
					__a							=   Image.open(requests.get(snake_case__      ,				stream=snake_case__						).raw						)

					return im
@torch.no_grad()
def 							lowerCAmelCase(	__lowerCamelCase      ,				__lowerCamelCase      ,				__lowerCamelCase						):
					__a							=   get_deta_config()

					# load original state dict
					if model_name == "deta-resnet-50":
										__a							=   'adet_checkpoint0011.pth'
					elif model_name == "deta-resnet-50-24-epochs":
										__a							=   'adet_2x_checkpoint0023.pth'
					else:
										raise ValueError(f'''Model name {model_name} not supported'''						)
					__a							=   hf_hub_download(repo_id='nielsr/deta-checkpoints'      ,				filename=snake_case__						)
					__a							=   torch.load(snake_case__      ,				map_location='cpu'						)['model']

					# rename keys
					__a							=   create_rename_keys(snake_case__						)
					for src, dest in rename_keys:
										rename_key(snake_case__      ,				snake_case__      ,				snake_case__						)
					read_in_decoder_q_k_v(snake_case__      ,				snake_case__						)

					# fix some prefixes
					for key in state_dict.copy().keys():
										if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
															__a							=   state_dict.pop(snake_case__						)
															__a							=   val
										if "input_proj" in key:
															__a							=   state_dict.pop(snake_case__						)
															__a							=   val
										if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
															__a							=   state_dict.pop(snake_case__						)
															__a							=   val

															 # finally, create HuggingFace model and load state dict
					__a							=   DetaForObjectDetection(snake_case__						)
					model.load_state_dict(snake_case__						)
					model.eval()

					__a							=   'cuda' if torch.cuda.is_available() else 'cpu'
					model.to(snake_case__						)

					# load image processor
					__a							=   DetaImageProcessor(format='coco_detection'						)

					# verify our conversion on image
					__a							=   prepare_img()
					__a							=   processor(images=snake_case__      ,				return_tensors='pt'						)
					__a							=   encoding['pixel_values']
					__a							=   model(pixel_values.to(snake_case__						)						)

					# verify logits
					if model_name == "deta-resnet-50":
										__a							=   torch.tensor(
										    [[-7.39_78, -2.54_06, -4.16_68], [-8.26_84, -3.99_33, -3.80_96], [-7.05_15, -3.79_73, -5.85_16]]						)
										__a							=   torch.tensor([[0.50_43, 0.49_73, 0.99_98], [0.25_42, 0.54_89, 0.47_48], [0.54_90, 0.27_65, 0.05_70]]						)
					elif model_name == "deta-resnet-50-24-epochs":
										__a							=   torch.tensor(
										    [[-7.16_88, -2.48_57, -4.86_69], [-7.86_30, -3.81_54, -4.26_74], [-7.27_30, -4.18_65, -5.53_23]]						)
										__a							=   torch.tensor([[0.50_21, 0.49_71, 0.99_94], [0.25_46, 0.54_86, 0.47_31], [0.16_86, 0.19_86, 0.21_42]]						)

					assert torch.allclose(outputs.logits[0, :3, :3]      ,				expected_logits.to(snake_case__						)      ,				atol=1e-4						)
					assert torch.allclose(outputs.pred_boxes[0, :3, :3]      ,				expected_boxes.to(snake_case__						)      ,				atol=1e-4						)
					print('Everything ok!'						)

					if pytorch_dump_folder_path:
										# Save model and processor
										logger.info(f'''Saving PyTorch model and processor to {pytorch_dump_folder_path}...'''						)
										Path(snake_case__						).mkdir(exist_ok=snake_case__						)
										model.save_pretrained(snake_case__						)
										processor.save_pretrained(snake_case__						)

					# Push to hub
					if push_to_hub:
										print('Pushing model and processor to hub...'						)
										model.push_to_hub(f'''jozhang97/{model_name}'''						)
										processor.push_to_hub(f'''jozhang97/{model_name}'''						)


if __name__ == "__main__":
	lowerCamelCase_				:		Dict       						=       argparse.ArgumentParser()

	parser.add_argument(
	    """--model_name""",
	    type=str,
	    default="""deta-resnet-50""",
	    choices=["""deta-resnet-50""", """deta-resnet-50-24-epochs"""],
	    help="""Name of the model you\'d like to convert.""",
	)
	parser.add_argument(
	    """--pytorch_dump_folder_path""",
	    default=None,
	    type=str,
	    help="""Path to the folder to output PyTorch model.""",
	)
	parser.add_argument(
	    """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
	)
	lowerCamelCase_				:		List[Any]       						=       parser.parse_args()
	convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


_UpperCamelCase			       =			namedtuple('''ExampleNamedTuple''', '''a b c''')



class     lowercase						(       unittest.TestCase					):



 '''simple docstring'''


 def   UpperCamelCase__							(self ) ->					Tuple:






  """simple docstring"""





  UpperCAmelCase__							=				torch.randn(5       ,       2 )
  UpperCAmelCase__							=				torch.device('cuda' ) if torch.cuda.is_available() else torch.device('cpu' )

  UpperCAmelCase__							=				send_to_device(__a       ,       __a )
  self.assertTrue(torch.equal(resulta.cpu()       ,       __a ) )

  UpperCAmelCase__							=				send_to_device((tensor, [tensor, tensor], 1)       ,       __a )
  self.assertIsInstance(__a       ,       __a )
  self.assertTrue(torch.equal(resulta[0].cpu()       ,       __a ) )
  self.assertIsInstance(resulta[1]       ,       __a )
  self.assertTrue(torch.equal(resulta[1][0].cpu()       ,       __a ) )
  self.assertTrue(torch.equal(resulta[1][1].cpu()       ,       __a ) )
  self.assertEqual(resulta[2]       ,       1 )

  UpperCAmelCase__							=				send_to_device({'a': tensor, 'b': [tensor, tensor], 'c': 1}       ,       __a )
  self.assertIsInstance(__a       ,       __a )
  self.assertTrue(torch.equal(resulta['a'].cpu()       ,       __a ) )
  self.assertIsInstance(resulta['b']       ,       __a )
  self.assertTrue(torch.equal(resulta['b'][0].cpu()       ,       __a ) )
  self.assertTrue(torch.equal(resulta['b'][1].cpu()       ,       __a ) )
  self.assertEqual(resulta['c']       ,       1 )

  UpperCAmelCase__							=				send_to_device(ExampleNamedTuple(a=__a       ,       b=[tensor, tensor]       ,       c=1 )       ,       __a )
  self.assertIsInstance(__a       ,       __a )
  self.assertTrue(torch.equal(resulta.a.cpu()       ,       __a ) )
  self.assertIsInstance(resulta.b       ,       __a )
  self.assertTrue(torch.equal(resulta.b[0].cpu()       ,       __a ) )
  self.assertTrue(torch.equal(resulta.b[1].cpu()       ,       __a ) )
  self.assertEqual(resulta.c       ,       1 )

  UpperCAmelCase__							=				send_to_device(UserDict({'a': tensor, 'b': [tensor, tensor], 'c': 1} )       ,       __a )
  self.assertIsInstance(__a       ,       __a )
  self.assertTrue(torch.equal(resulta['a'].cpu()       ,       __a ) )
  self.assertIsInstance(resulta['b']       ,       __a )
  self.assertTrue(torch.equal(resulta['b'][0].cpu()       ,       __a ) )
  self.assertTrue(torch.equal(resulta['b'][1].cpu()       ,       __a ) )
  self.assertEqual(resulta['c']       ,       1 )




 def   UpperCamelCase__							(self ) ->					List[str]:






  """simple docstring"""





  with self.assertRaises(__a ) as cm:
   UpperCAmelCase__							=				recursively_apply(torch.tensor       ,       (torch.tensor(1 ), 1)       ,       error_on_other_type=__a )
  self.assertEqual(
      str(cm.exception )       ,       'Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.'       ,       )




 def   UpperCamelCase__							(self ) ->					List[str]:






  """simple docstring"""





  UpperCAmelCase__							=				torch.tensor([1, 2, 3, 4, 5] )
  self.assertEqual(listify(__a )       ,       [1, 2, 3, 4, 5] )

  UpperCAmelCase__							=				torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]] )
  self.assertEqual(listify(__a )       ,       [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]] )

  UpperCAmelCase__							=				torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]] )
  self.assertEqual(
      listify(__a )       ,       [[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]] )




 def   UpperCamelCase__							(self ) ->					int:






  """simple docstring"""





  with patch_environment(aa=1       ,       BB=2 ):
   self.assertEqual(os.environ.get('AA' )       ,       '1' )
   self.assertEqual(os.environ.get('BB' )       ,       '2' )

  self.assertNotIn('AA'       ,       os.environ )
  self.assertNotIn('BB'       ,       os.environ )




 def   UpperCamelCase__							(self ) ->					Dict:






  """simple docstring"""





  UpperCAmelCase__							=				RegressionModel()
  UpperCAmelCase__							=				model.forward
  UpperCAmelCase__							=				convert_outputs_to_fpaa(model.forward )
  UpperCAmelCase__							=				extract_model_from_parallel(__a       ,       keep_fpaa_wrapper=__a )
  UpperCAmelCase__							=				pickle.dumps(__a )




 @require_cuda
 def   UpperCamelCase__							(self ) ->					Tuple:






  """simple docstring"""





  UpperCAmelCase__							=				RegressionModel()
  UpperCAmelCase__							=				model.forward
  UpperCAmelCase__							=				torch.cuda.amp.autocast(dtype=torch.floataa )(model.forward )
  UpperCAmelCase__							=				convert_outputs_to_fpaa(model.forward )
  UpperCAmelCase__							=				extract_model_from_parallel(__a       ,       keep_fpaa_wrapper=__a )
  UpperCAmelCase__							=				pickle.dumps(__a )




 @require_cuda
 @require_torch_min_version(version='2.0' )
 def   UpperCamelCase__							(self ) ->					List[str]:






  """simple docstring"""





  UpperCAmelCase__							=				RegressionModel()
  UpperCAmelCase__							=				model.forward
  UpperCAmelCase__							=				torch.cuda.amp.autocast(dtype=torch.floataa )(model.forward )
  UpperCAmelCase__							=				convert_outputs_to_fpaa(model.forward )
  UpperCAmelCase__							=				torch.compile(model.forward       ,       backend='inductor' )
  UpperCAmelCase__							=				torch.randn(4       ,       10 ).cuda()
  UpperCAmelCase__							=				model(__a )




 def   UpperCamelCase__							(self ) ->					Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__							=				RegressionModel()
  # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
  UpperCAmelCase__							=				torch.nn.parallel.DataParallel(__a )
  UpperCAmelCase__							=				extract_model_from_parallel(__a )

  self.assertEqual(__a       ,       __a )




 @require_torch_min_version(version='2.0' )
 def   UpperCamelCase__							(self ) ->					Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__							=				RegressionModel()
  UpperCAmelCase__							=				torch.compile(__a )

  # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
  UpperCAmelCase__							=				torch.nn.parallel.DataParallel(__a )
  UpperCAmelCase__							=				torch.compile(__a )
  UpperCAmelCase__							=				extract_model_from_parallel(__a )

  self.assertEqual(compiled_model._orig_mod       ,       compiled_model_unwrapped._orig_mod )




 def   UpperCamelCase__							(self ) ->					str:






  """simple docstring"""





  self.assertEqual(find_device([1, 'a', torch.tensor([1, 2, 3] )] )       ,       torch.device('cpu' ) )
  self.assertEqual(find_device({'a': 1, 'b': torch.tensor([1, 2, 3] )} )       ,       torch.device('cpu' ) )
  self.assertIsNone(find_device([1, 'a'] ) )



