


'''simple docstring'''





import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
		import torch

if is_vision_available():
		from PIL import Image

		from transformers import CLIPImageProcessor





class       _a							(       unittest.TestCase):




	def __init__(      self		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=7	,      _SCREAMING_SNAKE_CASE		:	List[str]=3	,      _SCREAMING_SNAKE_CASE		:	List[Any]=18	,      _SCREAMING_SNAKE_CASE		:	int=30	,      _SCREAMING_SNAKE_CASE		:	Dict=400	,      _SCREAMING_SNAKE_CASE		:	Any=True	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=None	,      _SCREAMING_SNAKE_CASE		:	int=True	,      _SCREAMING_SNAKE_CASE		:	int=None	,      _SCREAMING_SNAKE_CASE		:	List[Any]=True	,      _SCREAMING_SNAKE_CASE		:	Any=[0.4814_5466, 0.457_8275, 0.4082_1073]	,      _SCREAMING_SNAKE_CASE		:	Dict=[0.2686_2954, 0.2613_0258, 0.2757_7711]	,      _SCREAMING_SNAKE_CASE		:	List[Any]=True	,      )-> Any:
			lowerCAmelCase__     :    Tuple					  =      size if size is not None else {'''shortest_edge''': 20}
			lowerCAmelCase__     :    str					  =      crop_size if crop_size is not None else {'''height''': 18, '''width''': 18}
			lowerCAmelCase__     :    Dict					  =      parent
			lowerCAmelCase__     :    List[str]					  =      batch_size
			lowerCAmelCase__     :    List[str]					  =      num_channels
			lowerCAmelCase__     :    Optional[int]					  =      image_size
			lowerCAmelCase__     :    Optional[int]					  =      min_resolution
			lowerCAmelCase__     :    Union[str, Any]					  =      max_resolution
			lowerCAmelCase__     :    int					  =      do_resize
			lowerCAmelCase__     :    Any					  =      size
			lowerCAmelCase__     :    int					  =      do_center_crop
			lowerCAmelCase__     :    int					  =      crop_size
			lowerCAmelCase__     :    Optional[Any]					  =      do_normalize
			lowerCAmelCase__     :    Optional[int]					  =      image_mean
			lowerCAmelCase__     :    List[Any]					  =      image_std
			lowerCAmelCase__     :    str					  =      do_convert_rgb




	def 							UpperCAmelCase__(      self		:	Tuple	)-> int:
			return {
			    "do_resize": self.do_resize,
			    "size": self.size,
			    "do_center_crop": self.do_center_crop,
			    "crop_size": self.crop_size,
			    "do_normalize": self.do_normalize,
			    "image_mean": self.image_mean,
			    "image_std": self.image_std,
			    "do_convert_rgb": self.do_convert_rgb,
			}




	def 							UpperCAmelCase__(      self		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=False	,      _SCREAMING_SNAKE_CASE		:	List[Any]=False	,      _SCREAMING_SNAKE_CASE		:	int=False	)-> Optional[int]:
			assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

			if equal_resolution:
					lowerCAmelCase__     :    str					  =      []
					for i in range(self.batch_size	):
							image_inputs.append(
							    np.random.randint(
							        255	,      size=(self.num_channels, self.max_resolution, self.max_resolution)	,      dtype=np.uinta	)	)
			else:
					lowerCAmelCase__     :    Dict					  =      []
					for i in range(self.batch_size	):
							lowerCAmelCase__     :    Any					  =      np.random.choice(np.arange(self.min_resolution	,      self.max_resolution	)	,      2	)
							image_inputs.append(np.random.randint(255	,      size=(self.num_channels, width, height)	,      dtype=np.uinta	)	)

			if not numpify and not torchify:
					# PIL expects the channel dimension as last dimension
					lowerCAmelCase__     :    str					  =      [Image.fromarray(np.moveaxis(lowercase_	,      0	,      -1	)	) for x in image_inputs]

			if torchify:
					lowerCAmelCase__     :    Optional[Any]					  =      [torch.from_numpy(lowercase_	) for x in image_inputs]

			return image_inputs





@require_torch
@require_vision
class       _a							(       __UpperCAmelCase  ,		unittest.TestCase):
	_a						:						Dict   					=							CLIPImageProcessor if is_vision_available() else None




	def 							UpperCAmelCase__(      self		:	Any	)-> int:
			lowerCAmelCase__     :    Tuple					  =      CLIPImageProcessingTester(self	)




	@property
	def 							UpperCAmelCase__(      self		:	Union[str, Any]	)-> List[Any]:
			return self.image_processor_tester.prepare_image_processor_dict()




	def 							UpperCAmelCase__(      self		:	Any	)-> Tuple:
			lowerCAmelCase__     :    Tuple					  =      self.image_processing_class(**self.image_processor_dict	)
			self.assertTrue(hasattr(lowercase_	,      '''do_resize'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''size'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''do_center_crop'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''center_crop'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''do_normalize'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''image_mean'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''image_std'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''do_convert_rgb'''	)	)




	def 							UpperCAmelCase__(      self		:	Optional[Any]	)-> Any:
			lowerCAmelCase__     :    str					  =      self.image_processing_class.from_dict(self.image_processor_dict	)
			self.assertEqual(image_processor.size	,      {'''shortest_edge''': 20}	)
			self.assertEqual(image_processor.crop_size	,      {'''height''': 18, '''width''': 18}	)

			lowerCAmelCase__     :    Union[str, Any]					  =      self.image_processing_class.from_dict(self.image_processor_dict	,      size=42	,      crop_size=84	)
			self.assertEqual(image_processor.size	,      {'''shortest_edge''': 42}	)
			self.assertEqual(image_processor.crop_size	,      {'''height''': 84, '''width''': 84}	)




	def 							UpperCAmelCase__(      self		:	Optional[Any]	)-> Optional[int]:
			pass




	def 							UpperCAmelCase__(      self		:	str	)-> str:
			lowerCAmelCase__     :    Optional[Any]					  =      self.image_processing_class(**self.image_processor_dict	)
			# create random PIL images
			lowerCAmelCase__     :    List[str]					  =      self.image_processor_tester.prepare_inputs(equal_resolution=lowercase_	)
			for image in image_inputs:
					self.assertIsInstance(lowercase_	,      Image.Image	)

			# Test not batched input
			lowerCAmelCase__     :    Dict					  =      image_processing(image_inputs[0]	,      return_tensors='''pt'''	).pixel_values
			self.assertEqual(
			    encoded_images.shape	,      (
			        1,
			        self.image_processor_tester.num_channels,
			        self.image_processor_tester.crop_size['''height'''],
			        self.image_processor_tester.crop_size['''width'''],
			    )	,      )

			# Test batched
			lowerCAmelCase__     :    List[str]					  =      image_processing(lowercase_	,      return_tensors='''pt'''	).pixel_values
			self.assertEqual(
			    encoded_images.shape	,      (
			        self.image_processor_tester.batch_size,
			        self.image_processor_tester.num_channels,
			        self.image_processor_tester.crop_size['''height'''],
			        self.image_processor_tester.crop_size['''width'''],
			    )	,      )




	def 							UpperCAmelCase__(      self		:	int	)-> Tuple:
			lowerCAmelCase__     :    Any					  =      self.image_processing_class(**self.image_processor_dict	)
			# create random numpy tensors
			lowerCAmelCase__     :    Union[str, Any]					  =      self.image_processor_tester.prepare_inputs(equal_resolution=lowercase_	,      numpify=lowercase_	)
			for image in image_inputs:
					self.assertIsInstance(lowercase_	,      np.ndarray	)

			# Test not batched input
			lowerCAmelCase__     :    Dict					  =      image_processing(image_inputs[0]	,      return_tensors='''pt'''	).pixel_values
			self.assertEqual(
			    encoded_images.shape	,      (
			        1,
			        self.image_processor_tester.num_channels,
			        self.image_processor_tester.crop_size['''height'''],
			        self.image_processor_tester.crop_size['''width'''],
			    )	,      )

			# Test batched
			lowerCAmelCase__     :    Union[str, Any]					  =      image_processing(lowercase_	,      return_tensors='''pt'''	).pixel_values
			self.assertEqual(
			    encoded_images.shape	,      (
			        self.image_processor_tester.batch_size,
			        self.image_processor_tester.num_channels,
			        self.image_processor_tester.crop_size['''height'''],
			        self.image_processor_tester.crop_size['''width'''],
			    )	,      )




	def 							UpperCAmelCase__(      self		:	Optional[int]	)-> List[str]:
			lowerCAmelCase__     :    List[str]					  =      self.image_processing_class(**self.image_processor_dict	)
			# create random PyTorch tensors
			lowerCAmelCase__     :    Any					  =      self.image_processor_tester.prepare_inputs(equal_resolution=lowercase_	,      torchify=lowercase_	)
			for image in image_inputs:
					self.assertIsInstance(lowercase_	,      torch.Tensor	)

			# Test not batched input
			lowerCAmelCase__     :    Union[str, Any]					  =      image_processing(image_inputs[0]	,      return_tensors='''pt'''	).pixel_values
			self.assertEqual(
			    encoded_images.shape	,      (
			        1,
			        self.image_processor_tester.num_channels,
			        self.image_processor_tester.crop_size['''height'''],
			        self.image_processor_tester.crop_size['''width'''],
			    )	,      )

			# Test batched
			lowerCAmelCase__     :    Any					  =      image_processing(lowercase_	,      return_tensors='''pt'''	).pixel_values
			self.assertEqual(
			    encoded_images.shape	,      (
			        self.image_processor_tester.batch_size,
			        self.image_processor_tester.num_channels,
			        self.image_processor_tester.crop_size['''height'''],
			        self.image_processor_tester.crop_size['''width'''],
			    )	,      )





@require_torch
@require_vision
class       _a							(       __UpperCAmelCase  ,		unittest.TestCase):
	_a						:						Any   					=							CLIPImageProcessor if is_vision_available() else None




	def 							UpperCAmelCase__(      self		:	Optional[int]	)-> Dict:
			lowerCAmelCase__     :    Optional[int]					  =      CLIPImageProcessingTester(self	,      num_channels=4	)
			lowerCAmelCase__     :    List[str]					  =      3




	@property
	def 							UpperCAmelCase__(      self		:	Tuple	)-> Tuple:
			return self.image_processor_tester.prepare_image_processor_dict()




	def 							UpperCAmelCase__(      self		:	Optional[int]	)-> Optional[Any]:
			lowerCAmelCase__     :    List[str]					  =      self.image_processing_class(**self.image_processor_dict	)
			self.assertTrue(hasattr(lowercase_	,      '''do_resize'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''size'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''do_center_crop'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''center_crop'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''do_normalize'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''image_mean'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''image_std'''	)	)
			self.assertTrue(hasattr(lowercase_	,      '''do_convert_rgb'''	)	)




	def 							UpperCAmelCase__(      self		:	List[str]	)-> Dict:
			pass




	def 							UpperCAmelCase__(      self		:	Optional[Any]	)-> List[Any]:
			lowerCAmelCase__     :    List[Any]					  =      self.image_processing_class(**self.image_processor_dict	)
			# create random PIL images
			lowerCAmelCase__     :    Union[str, Any]					  =      self.image_processor_tester.prepare_inputs(equal_resolution=lowercase_	)
			for image in image_inputs:
					self.assertIsInstance(lowercase_	,      Image.Image	)

			# Test not batched input
			lowerCAmelCase__     :    Tuple					  =      image_processing(image_inputs[0]	,      return_tensors='''pt'''	).pixel_values
			self.assertEqual(
			    encoded_images.shape	,      (
			        1,
			        self.expected_encoded_image_num_channels,
			        self.image_processor_tester.crop_size['''height'''],
			        self.image_processor_tester.crop_size['''width'''],
			    )	,      )

			# Test batched
			lowerCAmelCase__     :    Tuple					  =      image_processing(lowercase_	,      return_tensors='''pt'''	).pixel_values
			self.assertEqual(
			    encoded_images.shape	,      (
			        self.image_processor_tester.batch_size,
			        self.expected_encoded_image_num_channels,
			        self.image_processor_tester.crop_size['''height'''],
			        self.image_processor_tester.crop_size['''width'''],
			    )	,      )








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


lowerCamelCase		=	logging.get_logger(__name__)





class       _a							(       _lowercase):
	_a						:						PyTorchBenchmarkArguments
	_a						:						PretrainedConfig
	_a						:						str   					=							"PyTorch"




	@property
	def 							UpperCAmelCase__(      self		:	Dict	)-> Any:
			return torch.__version__




	def 							UpperCAmelCase__(      self		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	)-> float:
			lowerCAmelCase__     :    Optional[Any]					  =      self._prepare_inference_func(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			return self._measure_speed(_inference	)




	def 							UpperCAmelCase__(      self		:	List[str]	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	)-> [Memory, Optional[MemorySummary]]:
			lowerCAmelCase__     :    List[Any]					  =      self._prepare_inference_func(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			return self._measure_memory(_inference	)




	def 							UpperCAmelCase__(      self		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	)-> float:
			lowerCAmelCase__     :    Any					  =      self._prepare_train_func(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			return self._measure_speed(_train	)




	def 							UpperCAmelCase__(      self		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	)-> [Memory, Optional[MemorySummary]]:
			lowerCAmelCase__     :    Tuple					  =      self._prepare_train_func(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			return self._measure_memory(_train	)




	def 							UpperCAmelCase__(      self		:	int	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	)-> Callable[[], None]:
			lowerCAmelCase__     :    int					  =      self.config_dict[model_name]

			if self.args.torchscript:
					lowerCAmelCase__     :    List[Any]					  =      True

			lowerCAmelCase__     :    List[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:
							lowerCAmelCase__     :    Dict					  =      config.architectures[0]
							lowerCAmelCase__     :    Optional[Any]					  =      __import__('''transformers'''	,      fromlist=[model_class]	)
							lowerCAmelCase__     :    Any					  =      getattr(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
							lowerCAmelCase__     :    List[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:
					lowerCAmelCase__     :    Union[str, Any]					  =      MODEL_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE	)

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

			# encoder-decoder has vocab size saved differently
			lowerCAmelCase__     :    str					  =      config.vocab_size if hasattr(_SCREAMING_SNAKE_CASE	,      '''vocab_size'''	) else config.encoder.vocab_size
			lowerCAmelCase__     :    Dict					  =      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():
							lowerCAmelCase__     :    int					  =      torch.jit.trace(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			else:
					lowerCAmelCase__     :    Optional[int]					  =      model

			def encoder_decoder_forward():
					with torch.no_grad():
							lowerCAmelCase__     :    str					  =      inference_model(_SCREAMING_SNAKE_CASE	,      decoder_input_ids=_SCREAMING_SNAKE_CASE	)
					return outputs

			def encoder_forward():
					with torch.no_grad():
							lowerCAmelCase__     :    List[Any]					  =      inference_model(_SCREAMING_SNAKE_CASE	)
					return outputs

			lowerCAmelCase__     :    Any					  =      encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
			return _forward




	def 							UpperCAmelCase__(      self		:	str	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	)-> Callable[[], None]:
			lowerCAmelCase__     :    Optional[int]					  =      self.config_dict[model_name]

			lowerCAmelCase__     :    List[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:
							lowerCAmelCase__     :    Optional[Any]					  =      config.architectures[0]
							lowerCAmelCase__     :    List[str]					  =      __import__('''transformers'''	,      fromlist=[model_class]	)
							lowerCAmelCase__     :    str					  =      getattr(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
							lowerCAmelCase__     :    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:
					lowerCAmelCase__     :    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:
					lowerCAmelCase__     :    Union[str, Any]					  =      model

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

			# encoder-decoder has vocab size saved differently
			lowerCAmelCase__     :    Union[str, Any]					  =      config.vocab_size if hasattr(_SCREAMING_SNAKE_CASE	,      '''vocab_size'''	) else config.encoder.vocab_size
			lowerCAmelCase__     :    List[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():
					lowerCAmelCase__     :    Optional[Any]					  =      train_model(_SCREAMING_SNAKE_CASE	,      labels=_SCREAMING_SNAKE_CASE	)[0]
					loss.backward()
					return loss

			def compute_loss_and_backprob_encoder_decoder():
					lowerCAmelCase__     :    List[str]					  =      train_model(_SCREAMING_SNAKE_CASE	,      decoder_input_ids=_SCREAMING_SNAKE_CASE	,      labels=_SCREAMING_SNAKE_CASE	)[0]
					loss.backward()
					return loss

			lowerCAmelCase__     :    List[Any]					  =      (
			    compute_loss_and_backprob_encoder_decoder
			    if config.is_encoder_decoder
			    else compute_loss_and_backprob_encoder
			)
			return _train




	def 							UpperCAmelCase__(      self		:	Union[str, Any]	,      _SCREAMING_SNAKE_CASE		:	List[str]	)-> float:
			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
					lowerCAmelCase__     :    Tuple					  =      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 							UpperCAmelCase__(      self		:	int	,      _SCREAMING_SNAKE_CASE		:	Callable[[], None]	)-> [Memory, MemorySummary]:
			try:
					if self.args.trace_memory_line_by_line:
							lowerCAmelCase__     :    int					  =      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.'''	)
									lowerCAmelCase__     :    Union[str, Any]					  =      '''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()
									lowerCAmelCase__     :    Optional[int]					  =      nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
									lowerCAmelCase__     :    Any					  =      nvml.nvmlDeviceGetMemoryInfo(_SCREAMING_SNAKE_CASE	)
									lowerCAmelCase__     :    int					  =      meminfo.used
									lowerCAmelCase__     :    Tuple					  =      Memory(_SCREAMING_SNAKE_CASE	)
									# shutdown nvml
									nvml.nvmlShutdown()
					else:
							# cpu
							lowerCAmelCase__     :    Union[str, Any]					  =      measure_peak_memory_cpu(_SCREAMING_SNAKE_CASE	)
							lowerCAmelCase__     :    Union[str, Any]					  =      Memory(_SCREAMING_SNAKE_CASE	) if isinstance(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	) else memory_bytes

					if self.args.trace_memory_line_by_line:
							lowerCAmelCase__     :    int					  =      stop_memory_tracing(_SCREAMING_SNAKE_CASE	)
					else:
							lowerCAmelCase__     :    str					  =      None

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




