
import argparse
import json
import os

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase__        =						logging.get_logger(__name__)
def 					lowerCAmelCase__     (    SCREAMING_SNAKE_CASE_	)							->						int:
			lowerCAmelCase__ : List[Any]   =						ConvNextVaConfig()

			if "atto" in checkpoint_url:
						lowerCAmelCase__ : List[str]   =						[2, 2, 6, 2]
						lowerCAmelCase__ : Optional[Any]   =						[40, 80, 160, 320]
			if "femto" in checkpoint_url:
						lowerCAmelCase__ : List[Any]   =						[2, 2, 6, 2]
						lowerCAmelCase__ : Dict   =						[48, 96, 192, 384]
			if "pico" in checkpoint_url:
						lowerCAmelCase__ : List[Any]   =						[2, 2, 6, 2]
						lowerCAmelCase__ : str   =						[64, 128, 256, 512]
			if "nano" in checkpoint_url:
						lowerCAmelCase__ : str   =						[2, 2, 8, 2]
						lowerCAmelCase__ : Optional[Any]   =						[80, 160, 320, 640]
			if "tiny" in checkpoint_url:
						lowerCAmelCase__ : Optional[Any]   =						[3, 3, 9, 3]
						lowerCAmelCase__ : int   =						[96, 192, 384, 768]
			if "base" in checkpoint_url:
						lowerCAmelCase__ : Dict   =						[3, 3, 27, 3]
						lowerCAmelCase__ : Any   =						[128, 256, 512, 1_024]
			if "large" in checkpoint_url:
						lowerCAmelCase__ : Optional[Any]   =						[3, 3, 27, 3]
						lowerCAmelCase__ : Tuple   =						[192, 384, 768, 1_536]
			if "huge" in checkpoint_url:
						lowerCAmelCase__ : List[Any]   =						[3, 3, 27, 3]
						lowerCAmelCase__ : Optional[int]   =						[352, 704, 1_408, 2_816]

			lowerCAmelCase__ : Tuple   =						1_000
			lowerCAmelCase__ : Optional[Any]   =						'imagenet-1k-id2label.json'
			lowerCAmelCase__ : Any   =						(1, 1_000)

			lowerCAmelCase__ : int   =						'huggingface/label-files'
			lowerCAmelCase__ : Optional[Any]   =						num_labels
			lowerCAmelCase__ : Tuple   =						json.load(open(hf_hub_download(_UpperCAmelCase			,				_UpperCAmelCase			,				repo_type='dataset'	)			,				'r'	)	)
			lowerCAmelCase__ : Optional[int]   =						{int(_UpperCAmelCase	): v for k, v in idalabel.items()}

			lowerCAmelCase__ : Dict   =						idalabel
			lowerCAmelCase__ : Any   =						{v: k for k, v in idalabel.items()}
			lowerCAmelCase__ : str   =						hidden_sizes
			lowerCAmelCase__ : Union[str, Any]   =						depths

			return config, expected_shape
def 					lowerCAmelCase__     (    SCREAMING_SNAKE_CASE_	)							->						Dict:
			if "downsample_layers.0.0" in name:
						lowerCAmelCase__ : str   =						name.replace('downsample_layers.0.0'			,				'embeddings.patch_embeddings'	)
			if "downsample_layers.0.1" in name:
						lowerCAmelCase__ : Tuple   =						name.replace('downsample_layers.0.1'			,				'embeddings.norm'	)  # we rename to layernorm later on
			if "downsample_layers.1.0" in name:
						lowerCAmelCase__ : Union[str, Any]   =						name.replace('downsample_layers.1.0'			,				'stages.1.downsampling_layer.0'	)
			if "downsample_layers.1.1" in name:
						lowerCAmelCase__ : Optional[Any]   =						name.replace('downsample_layers.1.1'			,				'stages.1.downsampling_layer.1'	)
			if "downsample_layers.2.0" in name:
						lowerCAmelCase__ : Tuple   =						name.replace('downsample_layers.2.0'			,				'stages.2.downsampling_layer.0'	)
			if "downsample_layers.2.1" in name:
						lowerCAmelCase__ : Optional[Any]   =						name.replace('downsample_layers.2.1'			,				'stages.2.downsampling_layer.1'	)
			if "downsample_layers.3.0" in name:
						lowerCAmelCase__ : List[Any]   =						name.replace('downsample_layers.3.0'			,				'stages.3.downsampling_layer.0'	)
			if "downsample_layers.3.1" in name:
						lowerCAmelCase__ : Optional[Any]   =						name.replace('downsample_layers.3.1'			,				'stages.3.downsampling_layer.1'	)
			if "stages" in name and "downsampling_layer" not in name:
						# stages.0.0. for instance should be renamed to stages.0.layers.0.
						lowerCAmelCase__ : Union[str, Any]   =						name[: len('stages.0'	)] + '.layers' + name[len('stages.0'	) :]
			if "gamma" in name:
						lowerCAmelCase__ : Dict   =						name.replace('gamma'			,				'weight'	)
			if "beta" in name:
						lowerCAmelCase__ : Optional[int]   =						name.replace('beta'			,				'bias'	)
			if "stages" in name:
						lowerCAmelCase__ : List[Any]   =						name.replace('stages'			,				'encoder.stages'	)
			if "norm" in name:
						lowerCAmelCase__ : Dict   =						name.replace('norm'			,				'layernorm'	)
			if "head" in name:
						lowerCAmelCase__ : str   =						name.replace('head'			,				'classifier'	)

			return name
def 					lowerCAmelCase__     (    )							->						Dict:
			lowerCAmelCase__ : Union[str, Any]   =						'http://images.cocodataset.org/val2017/000000039769.jpg'
			lowerCAmelCase__ : List[Any]   =						Image.open(requests.get(_UpperCAmelCase			,				stream=_UpperCAmelCase	).raw	)
			return im
def 					lowerCAmelCase__     (    SCREAMING_SNAKE_CASE_	)							->						List[str]:
			if "224" in checkpoint_url:
						lowerCAmelCase__ : Any   =						224
						lowerCAmelCase__ : Dict   =						224 / 256
			elif "384" in checkpoint_url:
						lowerCAmelCase__ : str   =						384
						lowerCAmelCase__ : List[Any]   =						None
			else:
						lowerCAmelCase__ : Optional[Any]   =						512
						lowerCAmelCase__ : Dict   =						None

			return ConvNextImageProcessor(
			    size=_UpperCAmelCase			,				crop_pct=_UpperCAmelCase			,				image_mean=[0.485, 0.456, 0.406]			,				image_std=[0.229, 0.224, 0.225]			,				resample=PILImageResampling.BICUBIC			,				)
@torch.no_grad()
def 					lowerCAmelCase__     (    SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_	)							->						Union[str, Any]:
			print('Downloading original model from checkpoint...'	)
			# define ConvNeXTV2 configuration based on URL
			lowerCAmelCase__ ,							lowerCAmelCase__ : Optional[int]   =						get_convnextva_config(_UpperCAmelCase	)
			# load original state_dict from URL
			lowerCAmelCase__ : Union[str, Any]   =						torch.hub.load_state_dict_from_url(_UpperCAmelCase	)['model']

			print('Converting model parameters...'	)
			# rename keys
			for key in state_dict.copy().keys():
						lowerCAmelCase__ : str   =						state_dict.pop(_UpperCAmelCase	)
						lowerCAmelCase__ : Any   =						val
			# add prefix to all keys expect classifier head
			for key in state_dict.copy().keys():
						lowerCAmelCase__ : str   =						state_dict.pop(_UpperCAmelCase	)
						if not key.startswith('classifier'	):
									lowerCAmelCase__ : List[str]   =						'convnextv2.' + key
						lowerCAmelCase__ : List[Any]   =						val

			# load HuggingFace model
			lowerCAmelCase__ : Optional[int]   =						ConvNextVaForImageClassification(_UpperCAmelCase	)
			model.load_state_dict(_UpperCAmelCase	)
			model.eval()

			# Check outputs on an image, prepared by ConvNextImageProcessor
			lowerCAmelCase__ : Optional[Any]   =						convert_preprocessor(_UpperCAmelCase	)
			lowerCAmelCase__ : str   =						preprocessor(images=prepare_img()			,				return_tensors='pt'	)
			lowerCAmelCase__ : Tuple   =						model(**_UpperCAmelCase	).logits

			# note: the logits below were obtained without center cropping
			if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
						lowerCAmelCase__ : Optional[Any]   =						torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
						lowerCAmelCase__ : List[Any]   =						torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
						lowerCAmelCase__ : Tuple   =						torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
						lowerCAmelCase__ : Tuple   =						torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
						lowerCAmelCase__ : Optional[int]   =						torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
						lowerCAmelCase__ : List[str]   =						torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
						lowerCAmelCase__ : Optional[Any]   =						torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
						lowerCAmelCase__ : Optional[int]   =						torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
						lowerCAmelCase__ : int   =						torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
						lowerCAmelCase__ : Optional[Any]   =						torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
						lowerCAmelCase__ : Any   =						torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
						lowerCAmelCase__ : Dict   =						torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
						lowerCAmelCase__ : Tuple   =						torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
						lowerCAmelCase__ : str   =						torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
						lowerCAmelCase__ : str   =						torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
						lowerCAmelCase__ : Optional[int]   =						torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
						lowerCAmelCase__ : int   =						torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309]	)
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
						lowerCAmelCase__ : Union[str, Any]   =						torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826]	)
			else:
						raise ValueError(F'''Unknown URL: {checkpoint_url}'''	)

			assert torch.allclose(logits[0, :5]			,				_UpperCAmelCase			,				atol=1e-3	)
			assert logits.shape == expected_shape
			print('Model outputs match the original results!'	)

			if save_model:
						print('Saving model to local...'	)
						# Create folder to save model
						if not os.path.isdir(_UpperCAmelCase	):
									os.mkdir(_UpperCAmelCase	)

						model.save_pretrained(_UpperCAmelCase	)
						preprocessor.save_pretrained(_UpperCAmelCase	)

			lowerCAmelCase__ : int   =						'convnextv2'
			if "atto" in checkpoint_url:
						model_name += "-atto"
			if "femto" in checkpoint_url:
						model_name += "-femto"
			if "pico" in checkpoint_url:
						model_name += "-pico"
			if "nano" in checkpoint_url:
						model_name += "-nano"
			elif "tiny" in checkpoint_url:
						model_name += "-tiny"
			elif "base" in checkpoint_url:
						model_name += "-base"
			elif "large" in checkpoint_url:
						model_name += "-large"
			elif "huge" in checkpoint_url:
						model_name += "-huge"
			if "22k" in checkpoint_url and "1k" not in checkpoint_url:
						model_name += "-22k"
			elif "22k" in checkpoint_url and "1k" in checkpoint_url:
						model_name += "-22k-1k"
			elif "1k" in checkpoint_url:
						model_name += "-1k"
			if "224" in checkpoint_url:
						model_name += "-224"
			elif "384" in checkpoint_url:
						model_name += "-384"
			elif "512" in checkpoint_url:
						model_name += "-512"

			if push_to_hub:
						print(F'''Pushing {model_name} to the hub...'''	)
						model.push_to_hub(_UpperCAmelCase	)
						preprocessor.push_to_hub(_UpperCAmelCase	)


if __name__ == "__main__":
					lowerCamelCase__        =						argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    """--checkpoint_url""",
					    default="""https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt""",
					    type=str,
					    help="""URL of the original ConvNeXTV2 checkpoint you\'d like to convert.""",
					)
					parser.add_argument(
					    """--pytorch_dump_folder_path""",
					    default="""model""",
					    type=str,
					    help="""Path to the output PyTorch model directory.""",
					)
					parser.add_argument("""--save_model""", action="""store_true""", help="""Save model to local""")
					parser.add_argument("""--push_to_hub""", action="""store_true""", help="""Push model and image preprocessor to the hub""")

					lowerCamelCase__        =						parser.parse_args()
					convert_convnextva_checkpoint(
					    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, 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__    :      List[Any]       					=namedtuple('''ExampleNamedTuple''', '''a b c''')

class    __A  (    unittest.TestCase     ):


				def _snake_case							(		self				):
								lowerCamelCase									=torch.randn(5 ,			2				)
								lowerCamelCase									=torch.device("""cuda"""				) if torch.cuda.is_available() else torch.device("""cpu"""				)

								lowerCamelCase									=send_to_device(UpperCAmelCase_ ,			UpperCAmelCase_				)
								self.assertTrue(torch.equal(resulta.cpu() ,			UpperCAmelCase_				)				)

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

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

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

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


				def _snake_case							(		self				):
								with self.assertRaises(UpperCAmelCase_				) as cm:
												lowerCamelCase									=recursively_apply(torch.tensor ,			(torch.tensor(1				), 1) ,			error_on_other_type=UpperCAmelCase_				)
								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 _snake_case							(		self				):
								lowerCamelCase									=torch.tensor([1, 2, 3, 4, 5]				)
								self.assertEqual(listify(UpperCAmelCase_				) ,			[1, 2, 3, 4, 5]				)

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

								lowerCamelCase									=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(UpperCAmelCase_				) ,			[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]				)


				def _snake_case							(		self				):
								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 _snake_case							(		self				):
								lowerCamelCase									=RegressionModel()
								lowerCamelCase									=model.forward
								lowerCamelCase									=convert_outputs_to_fpaa(model.forward				)
								lowerCamelCase									=extract_model_from_parallel(UpperCAmelCase_ ,			keep_fpaa_wrapper=UpperCAmelCase_				)
								lowerCamelCase									=pickle.dumps(UpperCAmelCase_				)


				@require_cuda
				def _snake_case							(		self				):
								lowerCamelCase									=RegressionModel()
								lowerCamelCase									=model.forward
								lowerCamelCase									=torch.cuda.amp.autocast(dtype=torch.floataa				)(model.forward				)
								lowerCamelCase									=convert_outputs_to_fpaa(model.forward				)
								lowerCamelCase									=extract_model_from_parallel(UpperCAmelCase_ ,			keep_fpaa_wrapper=UpperCAmelCase_				)
								lowerCamelCase									=pickle.dumps(UpperCAmelCase_				)


				@require_cuda
				@require_torch_min_version(version="""2.0"""				)
				def _snake_case							(		self				):
								lowerCamelCase									=RegressionModel()
								lowerCamelCase									=model.forward
								lowerCamelCase									=torch.cuda.amp.autocast(dtype=torch.floataa				)(model.forward				)
								lowerCamelCase									=convert_outputs_to_fpaa(model.forward				)
								lowerCamelCase									=torch.compile(model.forward ,			backend="""inductor"""				)
								lowerCamelCase									=torch.randn(4 ,			10				).cuda()
								lowerCamelCase									=model(UpperCAmelCase_				)


				def _snake_case							(		self				):
								lowerCamelCase									=RegressionModel()
								# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
								lowerCamelCase									=torch.nn.parallel.DataParallel(UpperCAmelCase_				)
								lowerCamelCase									=extract_model_from_parallel(UpperCAmelCase_				)

								self.assertEqual(UpperCAmelCase_ ,			UpperCAmelCase_				)


				@require_torch_min_version(version="""2.0"""				)
				def _snake_case							(		self				):
								lowerCamelCase									=RegressionModel()
								lowerCamelCase									=torch.compile(UpperCAmelCase_				)

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

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


				def _snake_case							(		self				):
								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"""]				)				)





