



"""simple docstring"""


from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import ConvNextConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers import TFConvNextForImageClassification, TFConvNextModel


if is_vision_available():
				from PIL import Image

				from transformers import ConvNextImageProcessor







class  A_     :



						"""simple docstring"""






						def __init__(   self					:str  ,    lowercase_					:Optional[Any]  ,    lowercase_					:Tuple=13  ,    lowercase_					:Any=32  ,    lowercase_					:Dict=3  ,    lowercase_					:Union[str, Any]=4  ,    lowercase_					:Union[str, Any]=[10, 20, 30, 40]  ,    lowercase_					:Dict=[2, 2, 3, 2]  ,    lowercase_					:int=True  ,    lowercase_					:List[str]=True  ,    lowercase_					:Optional[int]=37  ,    lowercase_					:Tuple="gelu"  ,    lowercase_					:int=10  ,    lowercase_					:Union[str, Any]=0.02  ,    lowercase_					:Dict=3  ,    lowercase_					:Tuple=None  ,    )							->     int:
												UpperCAmelCase			 =					parent
												UpperCAmelCase			 =					batch_size
												UpperCAmelCase			 =					image_size
												UpperCAmelCase			 =					num_channels
												UpperCAmelCase			 =					num_stages
												UpperCAmelCase			 =					hidden_sizes
												UpperCAmelCase			 =					depths
												UpperCAmelCase			 =					is_training
												UpperCAmelCase			 =					use_labels
												UpperCAmelCase			 =					intermediate_size
												UpperCAmelCase			 =					hidden_act
												UpperCAmelCase			 =					type_sequence_label_size
												UpperCAmelCase			 =					initializer_range
												UpperCAmelCase			 =					scope






						def       UpperCAmelCase__	(   self					:str   )							->     Optional[int]:
												UpperCAmelCase			 =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

												UpperCAmelCase			 =					None
												if self.use_labels:
																		UpperCAmelCase			 =					ids_tensor([self.batch_size]  ,    self.type_sequence_label_size   )

												UpperCAmelCase			 =					self.get_config()

												return config, pixel_values, labels






						def       UpperCAmelCase__	(   self					:Union[str, Any]   )							->     str:
												return ConvNextConfig(
												    num_channels=self.num_channels  ,    hidden_sizes=self.hidden_sizes  ,    depths=self.depths  ,    num_stages=self.num_stages  ,    hidden_act=self.hidden_act  ,    is_decoder=lowercase_  ,    initializer_range=self.initializer_range  ,    )






						def       UpperCAmelCase__	(   self					:Dict  ,    lowercase_					:Dict  ,    lowercase_					:Tuple  ,    lowercase_					:Optional[Any]   )							->     Any:
												UpperCAmelCase			 =					TFConvNextModel(config=lowercase_   )
												UpperCAmelCase			 =					model(lowercase_  ,    training=lowercase_   )
												# expected last hidden states: B, C, H // 32, W // 32
												self.parent.assertEqual(
												    result.last_hidden_state.shape  ,    (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)  ,    )






						def       UpperCAmelCase__	(   self					:Optional[Any]  ,    lowercase_					:Optional[int]  ,    lowercase_					:List[Any]  ,    lowercase_					:Optional[Any]   )							->     Optional[int]:
												UpperCAmelCase			 =					self.type_sequence_label_size
												UpperCAmelCase			 =					TFConvNextForImageClassification(lowercase_   )
												UpperCAmelCase			 =					model(lowercase_  ,    labels=lowercase_  ,    training=lowercase_   )
												self.parent.assertEqual(result.logits.shape  ,    (self.batch_size, self.type_sequence_label_size)   )






						def       UpperCAmelCase__	(   self					:int   )							->     Dict:
												UpperCAmelCase			 =					self.prepare_config_and_inputs()
												UpperCAmelCase ,       UpperCAmelCase ,       UpperCAmelCase			 =					config_and_inputs
												UpperCAmelCase			 =					{'pixel_values': pixel_values}
												return config, inputs_dict







@require_tf
class  A_     (	SCREAMING_SNAKE_CASE_      ,  SCREAMING_SNAKE_CASE_      ,  unittest.TestCase							):



						"""simple docstring"""
						__UpperCamelCase													=				(TFConvNextModel, TFConvNextForImageClassification) if is_tf_available() else ()
						__UpperCamelCase													=				(
						    {"""feature-extraction""": TFConvNextModel, """image-classification""": TFConvNextForImageClassification}
						    if is_tf_available()
						    else {}
						)

						__UpperCamelCase													=				False
						__UpperCamelCase													=				False
						__UpperCamelCase													=				False
						__UpperCamelCase													=				False
						__UpperCamelCase													=				False






						def       UpperCAmelCase__	(   self					:List[Any]   )							->     str:
												UpperCAmelCase			 =					TFConvNextModelTester(self   )
												UpperCAmelCase			 =					ConfigTester(
												    self  ,    config_class=lowercase_  ,    has_text_modality=lowercase_  ,    hidden_size=37  ,    )






						@unittest.skip(reason='ConvNext does not use inputs_embeds'   )
						def       UpperCAmelCase__	(   self					:Any   )							->     Optional[int]:
												pass






						@unittest.skipIf(
						    not is_tf_available() or len(tf.config.list_physical_devices('GPU'   )   ) == 0  ,    reason='TF does not support backprop for grouped convolutions on CPU.'  ,    )
						@slow
						def       UpperCAmelCase__	(   self					:Any   )							->     Optional[int]:
												super().test_keras_fit()






						@unittest.skip(reason='ConvNext does not support input and output embeddings'   )
						def       UpperCAmelCase__	(   self					:List[str]   )							->     List[Any]:
												pass






						def       UpperCAmelCase__	(   self					:str   )							->     Optional[int]:
												UpperCAmelCase ,       UpperCAmelCase			 =					self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																		UpperCAmelCase			 =					model_class(lowercase_   )
																		UpperCAmelCase			 =					inspect.signature(model.call   )
																		# signature.parameters is an OrderedDict => so arg_names order is deterministic
																		UpperCAmelCase			 =					[*signature.parameters.keys()]

																		UpperCAmelCase			 =					['pixel_values']
																		self.assertListEqual(arg_names[:1]  ,    lowercase_   )






						def       UpperCAmelCase__	(   self					:List[Any]   )							->     Dict:
												UpperCAmelCase			 =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*lowercase_   )






						@unittest.skipIf(
						    not is_tf_available() or len(tf.config.list_physical_devices('GPU'   )   ) == 0  ,    reason='TF does not support backprop for grouped convolutions on CPU.'  ,    )
						def       UpperCAmelCase__	(   self					:List[str]   )							->     str:
												super().test_dataset_conversion()






						def       UpperCAmelCase__	(   self					:Dict   )							->     Any:
												def check_hidden_states_output(lowercase_					:Dict  ,    lowercase_					:str  ,    lowercase_					:Optional[Any]   ):
																		UpperCAmelCase			 =					model_class(lowercase_   )

																		UpperCAmelCase			 =					model(**self._prepare_for_class(lowercase_  ,    lowercase_   )   )
																		UpperCAmelCase			 =					outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																		UpperCAmelCase			 =					self.model_tester.num_stages
																		self.assertEqual(len(lowercase_   )  ,    expected_num_stages + 1   )

																		# ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
																		self.assertListEqual(
																		    list(hidden_states[0].shape[-2:]   )  ,    [self.model_tester.image_size // 4, self.model_tester.image_size // 4]  ,    )

												UpperCAmelCase ,       UpperCAmelCase			 =					self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																		UpperCAmelCase			 =					True
																		check_hidden_states_output(lowercase_  ,    lowercase_  ,    lowercase_   )

																		# check that output_hidden_states also work using config
																		del inputs_dict["output_hidden_states"]
																		UpperCAmelCase			 =					True

																		check_hidden_states_output(lowercase_  ,    lowercase_  ,    lowercase_   )






						def       UpperCAmelCase__	(   self					:Optional[int]   )							->     Optional[int]:
												UpperCAmelCase ,       UpperCAmelCase			 =					self.model_tester.prepare_config_and_inputs_for_common()

												def check_equivalence(lowercase_					:Union[str, Any]  ,    lowercase_					:str  ,    lowercase_					:Optional[Any]  ,    lowercase_					:Tuple={}   ):
																		UpperCAmelCase			 =					model(lowercase_  ,    return_dict=lowercase_  ,    **lowercase_   )
																		UpperCAmelCase			 =					model(lowercase_  ,    return_dict=lowercase_  ,    **lowercase_   ).to_tuple()

																		def recursive_check(lowercase_					:Optional[int]  ,    lowercase_					:Union[str, Any]   ):
																								if isinstance(lowercase_  ,    (List, Tuple)   ):
																														for tuple_iterable_value, dict_iterable_value in zip(lowercase_  ,    lowercase_   ):
																																				recursive_check(lowercase_  ,    lowercase_   )
																								elif tuple_object is None:
																														return
																								else:
																														self.assertTrue(
																														    all(tf.equal(lowercase_  ,    lowercase_   )   )  ,    msg=(
																														        'Tuple and dict output are not equal. Difference:'
																														        f""" {tf.math.reduce_max(tf.abs(tuple_object - dict_object   )   )}"""
																														    )  ,    )

																								recursive_check(lowercase_  ,    lowercase_   )

												for model_class in self.all_model_classes:
																		UpperCAmelCase			 =					model_class(lowercase_   )

																		UpperCAmelCase			 =					self._prepare_for_class(lowercase_  ,    lowercase_   )
																		UpperCAmelCase			 =					self._prepare_for_class(lowercase_  ,    lowercase_   )
																		check_equivalence(lowercase_  ,    lowercase_  ,    lowercase_   )

																		UpperCAmelCase			 =					self._prepare_for_class(lowercase_  ,    lowercase_  ,    return_labels=lowercase_   )
																		UpperCAmelCase			 =					self._prepare_for_class(lowercase_  ,    lowercase_  ,    return_labels=lowercase_   )
																		check_equivalence(lowercase_  ,    lowercase_  ,    lowercase_   )

																		UpperCAmelCase			 =					self._prepare_for_class(lowercase_  ,    lowercase_   )
																		UpperCAmelCase			 =					self._prepare_for_class(lowercase_  ,    lowercase_   )
																		check_equivalence(lowercase_  ,    lowercase_  ,    lowercase_  ,    {'output_hidden_states': True}   )

																		UpperCAmelCase			 =					self._prepare_for_class(lowercase_  ,    lowercase_  ,    return_labels=lowercase_   )
																		UpperCAmelCase			 =					self._prepare_for_class(lowercase_  ,    lowercase_  ,    return_labels=lowercase_   )
																		check_equivalence(lowercase_  ,    lowercase_  ,    lowercase_  ,    {'output_hidden_states': True}   )






						def       UpperCAmelCase__	(   self					:Tuple   )							->     List[str]:
												UpperCAmelCase			 =					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_image_classification(*lowercase_   )






						@slow
						def       UpperCAmelCase__	(   self					:Tuple   )							->     Union[str, Any]:
												UpperCAmelCase			 =					TFConvNextModel.from_pretrained('facebook/convnext-tiny-224'   )
												self.assertIsNotNone(lowercase_   )


def 					_lowerCAmelCase (  ):
						UpperCAmelCase			 =					Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'		)
						return image







@require_tf
@require_vision
class  A_     (	unittest.TestCase							):



						"""simple docstring"""






						@cached_property
						def       UpperCAmelCase__	(   self					:Union[str, Any]   )							->     Dict:
												return ConvNextImageProcessor.from_pretrained('facebook/convnext-tiny-224'   ) if is_vision_available() else None






						@slow
						def       UpperCAmelCase__	(   self					:Any   )							->     Optional[Any]:
												UpperCAmelCase			 =					TFConvNextForImageClassification.from_pretrained('facebook/convnext-tiny-224'   )

												UpperCAmelCase			 =					self.default_image_processor
												UpperCAmelCase			 =					prepare_img()
												UpperCAmelCase			 =					image_processor(images=lowercase_  ,    return_tensors='tf'   )

												# forward pass
												UpperCAmelCase			 =					model(**lowercase_   )

												# verify the logits
												UpperCAmelCase			 =					tf.TensorShape((1, 10_00)   )
												self.assertEqual(outputs.logits.shape  ,    lowercase_   )

												UpperCAmelCase			 =					tf.constant([-0.0260, -0.4739, 0.1911]   )

												tf.debugging.assert_near(outputs.logits[0, :3]  ,    lowercase_  ,    atol=1E-4   )











'''simple docstring'''


import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
     import torch

     from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
     import pyanvml.pyanvml as nvml


UpperCamelCase			     =					logging.get_logger(__name__)






class 						lowerCAmelCase_	(       UpperCAmelCase_							):







  '''simple docstring'''

  UpperCamelCase_			:      PyTorchBenchmarkArguments
  UpperCamelCase_			:      PretrainedConfig
  UpperCamelCase_			:      str										=				"PyTorch"



  @property
  def 					_snake_case     ( self				:							int	)			->		List[str]:







       '''simple docstring'''



       return torch.__version__



  def 					_snake_case     ( self				:							Dict					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		float:







       '''simple docstring'''



       A:		List[str]          =		self._prepare_inference_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_speed(_inference	)



  def 					_snake_case     ( self				:							Dict					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		[Memory, Optional[MemorySummary]]:







       '''simple docstring'''



       A:		Any          =		self._prepare_inference_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_memory(_inference	)



  def 					_snake_case     ( self				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		float:







       '''simple docstring'''



       A:		List[Any]          =		self._prepare_train_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_speed(_train	)



  def 					_snake_case     ( self				:							int					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		[Memory, Optional[MemorySummary]]:







       '''simple docstring'''



       A:		List[Any]          =		self._prepare_train_func(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       return self._measure_memory(_train	)



  def 					_snake_case     ( self				:							Union[str, Any]					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		Callable[[], None]:







       '''simple docstring'''



       A:		int          =		self.config_dict[model_name]

       if self.args.torchscript:
            A:		int          =		True

       A:		Any          =		(
           hasattr(SCREAMING_SNAKE_CASE_					,  '''architectures'''	)
           and isinstance(config.architectures					,  SCREAMING_SNAKE_CASE_	)
           and len(config.architectures	) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                 A:		Optional[int]          =		config.architectures[0]
                 A:		Optional[int]          =		__import__('''transformers'''					,  fromlist=[model_class]	)
                 A:		Dict          =		getattr(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
                 A:		Optional[Any]          =		model_cls(SCREAMING_SNAKE_CASE_	)
            except ImportError:
                 raise ImportError(
                     f"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                     ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''	)
       else:
            A:		Union[str, Any]          =		MODEL_MAPPING[config.__class__](SCREAMING_SNAKE_CASE_	)

       model.eval()
       model.to(self.args.device	)

       # encoder-decoder has vocab size saved differently
       A:		Tuple          =		config.vocab_size if hasattr(SCREAMING_SNAKE_CASE_					,  '''vocab_size'''	) else config.encoder.vocab_size
       A:		List[Any]          =		torch.randint(SCREAMING_SNAKE_CASE_					,  (batch_size, sequence_length)					,  dtype=torch.long					,  device=self.args.device	)

       if self.args.fpaa:
            logger.info('''Running training in Mixed Precision...'''	)
            if not self.args.is_gpu:
                 raise ValueError('''Mixed precision is possible only for GPU.'''	)
            # amp seems to have memory leaks so that memory usage
            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
            model.half()

       if self.args.torchscript:
            with torch.no_grad():
                 A:		Any          =		torch.jit.trace(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       else:
            A:		Dict          =		model

       def encoder_decoder_forward():
            with torch.no_grad():
                 A:		List[str]          =		inference_model(SCREAMING_SNAKE_CASE_					,  decoder_input_ids=SCREAMING_SNAKE_CASE_	)
            return outputs

       def encoder_forward():
            with torch.no_grad():
                 A:		Union[str, Any]          =		inference_model(SCREAMING_SNAKE_CASE_	)
            return outputs

       A:		List[Any]          =		encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
       return _forward



  def 					_snake_case     ( self				:							Tuple					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							int	)			->		Callable[[], None]:







       '''simple docstring'''



       A:		int          =		self.config_dict[model_name]

       A:		str          =		(
           hasattr(SCREAMING_SNAKE_CASE_					,  '''architectures'''	)
           and isinstance(config.architectures					,  SCREAMING_SNAKE_CASE_	)
           and len(config.architectures	) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                 A:		List[str]          =		config.architectures[0]
                 A:		List[str]          =		__import__('''transformers'''					,  fromlist=[model_class]	)
                 A:		Optional[int]          =		getattr(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
                 A:		Union[str, Any]          =		model_cls(SCREAMING_SNAKE_CASE_	)
            except ImportError:
                 raise ImportError(
                     f"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                     ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''	)
       else:
            A:		Union[str, Any]          =		MODEL_WITH_LM_HEAD_MAPPING[config.__class__](SCREAMING_SNAKE_CASE_	)

       if self.args.torchscript:
            raise NotImplementedError('''Training for torchscript is currently not implemented'''	)
       else:
            A:		str          =		model

       model.train()
       model.to(self.args.device	)

       # encoder-decoder has vocab size saved differently
       A:		Tuple          =		config.vocab_size if hasattr(SCREAMING_SNAKE_CASE_					,  '''vocab_size'''	) else config.encoder.vocab_size
       A:		str          =		torch.randint(SCREAMING_SNAKE_CASE_					,  (batch_size, sequence_length)					,  dtype=torch.long					,  device=self.args.device	)

       if self.args.fpaa:
            logger.info('''Running training in Mixed Precision...'''	)
            if not self.args.is_gpu:
                 raise ValueError('''Mixed precision is possible only for GPU.'''	)

            # amp seems to have memory leaks so that memory usage
            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
            model.half()

       def compute_loss_and_backprob_encoder():
            A:		str          =		train_model(SCREAMING_SNAKE_CASE_					,  labels=SCREAMING_SNAKE_CASE_	)[0]
            loss.backward()
            return loss

       def compute_loss_and_backprob_encoder_decoder():
            A:		Dict          =		train_model(SCREAMING_SNAKE_CASE_					,  decoder_input_ids=SCREAMING_SNAKE_CASE_					,  labels=SCREAMING_SNAKE_CASE_	)[0]
            loss.backward()
            return loss

       A:		str          =		(
           compute_loss_and_backprob_encoder_decoder
           if config.is_encoder_decoder
           else compute_loss_and_backprob_encoder
       )
       return _train



  def 					_snake_case     ( self				:							Any					,  SCREAMING_SNAKE_CASE_				:							Any	)			->		float:







       '''simple docstring'''



       try:
            if self.args.is_tpu or self.args.torchscript:
                 # run additional 10 times to stabilize compilation for tpu and torchscript
                 logger.info('''Do inference on TPU or torchscript. Running model 5 times to stabilize compilation'''	)
                 timeit.repeat(
                     SCREAMING_SNAKE_CASE_					,  repeat=1					,  number=5					,  )

            # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
            A:		Any          =		timeit.repeat(
                SCREAMING_SNAKE_CASE_					,  repeat=self.args.repeat					,  number=10					,  )

            if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
                 import torch_xla.debug.metrics as met

                 self.print_fn(met.metrics_report()	)

            return min(SCREAMING_SNAKE_CASE_	) / 10.0
       except RuntimeError as e:
            self.print_fn(f"""Doesn't fit on GPU. {e}"""	)
            return "N/A"




  def 					_snake_case     ( self				:							Optional[int]					,  SCREAMING_SNAKE_CASE_				:							Callable[[], None]	)			->		[Memory, MemorySummary]:







       '''simple docstring'''



       try:
            if self.args.trace_memory_line_by_line:
                 A:		List[Any]          =		start_memory_tracing('''transformers'''	)

            if self.args.is_tpu:
                 # tpu
                 raise NotImplementedError(
                     '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'''
                     ''' `--no-memory` or `args.memory=False`'''	)
            elif self.args.is_gpu:
                 if not is_pyanvml_available():
                      logger.warning(
                          '''py3nvml not installed, we won\'t log GPU memory usage. '''
                          '''Install py3nvml (pip install py3nvml) to log information about GPU.'''	)
                      A:		Optional[int]          =		'''N/A'''
                 else:
                      logger.info(
                          '''Measuring total GPU usage on GPU device. Make sure to not have additional processes running'''
                          ''' on the same GPU.'''	)
                      # init nvml
                      nvml.nvmlInit()
                      func()
                      A:		Union[str, Any]          =		nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
                      A:		Union[str, Any]          =		nvml.nvmlDeviceGetMemoryInfo(SCREAMING_SNAKE_CASE_	)
                      A:		Tuple          =		meminfo.used
                      A:		Tuple          =		Memory(SCREAMING_SNAKE_CASE_	)
                      # shutdown nvml
                      nvml.nvmlShutdown()
            else:
                 # cpu
                 A:		int          =		measure_peak_memory_cpu(SCREAMING_SNAKE_CASE_	)
                 A:		str          =		Memory(SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	) else memory_bytes

            if self.args.trace_memory_line_by_line:
                 A:		Dict          =		stop_memory_tracing(SCREAMING_SNAKE_CASE_	)
            else:
                 A:		List[str]          =		None

            return memory, summary
       except RuntimeError as e:
            self.print_fn(f"""Doesn't fit on GPU. {e}"""	)
            return "N/A", None






