




import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__SCREAMING_SNAKE_CASE			    =logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE			    ={
    """asapp/sew-d-tiny-100k""": """https://huggingface.co/asapp/sew-d-tiny-100k/resolve/main/config.json""",
    # See all SEW-D models at https://huggingface.co/models?filter=sew-d
}






class 			__magic_name__    (	__UpperCAmelCase):
				'''simple docstring'''







				SCREAMING_SNAKE_CASE__       :				str						   =       """sew-d"""


				def __init__(				self: List[str]		,       _lowerCamelCase: Dict=32		,       _lowerCamelCase: List[Any]=7_68		,       _lowerCamelCase: List[str]=12		,       _lowerCamelCase: Tuple=12		,       _lowerCamelCase: Optional[int]=30_72		,       _lowerCamelCase: Any=2		,       _lowerCamelCase: Union[str, Any]=5_12		,       _lowerCamelCase: str=2_56		,       _lowerCamelCase: List[Any]=True		,       _lowerCamelCase: Optional[int]=True		,       _lowerCamelCase: Tuple=("p2c", "c2p")		,       _lowerCamelCase: str="layer_norm"		,       _lowerCamelCase: int="gelu_python"		,       _lowerCamelCase: Union[str, Any]=0.1		,       _lowerCamelCase: Optional[Any]=0.1		,       _lowerCamelCase: int=0.1		,       _lowerCamelCase: Union[str, Any]=0.0		,       _lowerCamelCase: int=0.1		,       _lowerCamelCase: int=0.02		,       _lowerCamelCase: List[Any]=1E-7		,       _lowerCamelCase: int=1E-5		,       _lowerCamelCase: List[str]="group"		,       _lowerCamelCase: Union[str, Any]="gelu"		,       _lowerCamelCase: Any=(64, 1_28, 1_28, 1_28, 1_28, 2_56, 2_56, 2_56, 2_56, 5_12, 5_12, 5_12, 5_12)		,       _lowerCamelCase: str=(5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1)		,       _lowerCamelCase: int=(10, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1)		,       _lowerCamelCase: Dict=False		,       _lowerCamelCase: str=1_28		,       _lowerCamelCase: Optional[int]=16		,       _lowerCamelCase: int=True		,       _lowerCamelCase: Any=0.05		,       _lowerCamelCase: Optional[int]=10		,       _lowerCamelCase: int=2		,       _lowerCamelCase: Any=0.0		,       _lowerCamelCase: Optional[Any]=10		,       _lowerCamelCase: str=0		,       _lowerCamelCase: str="mean"		,       _lowerCamelCase: List[str]=False		,       _lowerCamelCase: List[Any]=False		,       _lowerCamelCase: List[str]=2_56		,       _lowerCamelCase: Tuple=0		,       _lowerCamelCase: Union[str, Any]=1		,       _lowerCamelCase: List[str]=2		,       **_lowerCamelCase: Tuple		,       ):
											super().__init__(**_lowercase		,       pad_token_id=_lowercase		,       bos_token_id=_lowercase		,       eos_token_id=_lowercase		)
											SCREAMING_SNAKE_CASE_ 						=     hidden_size
											SCREAMING_SNAKE_CASE_ 						=     feat_extract_norm
											SCREAMING_SNAKE_CASE_ 						=     feat_extract_activation
											SCREAMING_SNAKE_CASE_ 						=     list(_lowercase		)
											SCREAMING_SNAKE_CASE_ 						=     list(_lowercase		)
											SCREAMING_SNAKE_CASE_ 						=     list(_lowercase		)
											SCREAMING_SNAKE_CASE_ 						=     conv_bias
											SCREAMING_SNAKE_CASE_ 						=     num_conv_pos_embeddings
											SCREAMING_SNAKE_CASE_ 						=     num_conv_pos_embedding_groups
											SCREAMING_SNAKE_CASE_ 						=     len(self.conv_dim		)
											SCREAMING_SNAKE_CASE_ 						=     num_hidden_layers
											SCREAMING_SNAKE_CASE_ 						=     intermediate_size
											SCREAMING_SNAKE_CASE_ 						=     squeeze_factor
											SCREAMING_SNAKE_CASE_ 						=     max_position_embeddings
											SCREAMING_SNAKE_CASE_ 						=     position_buckets
											SCREAMING_SNAKE_CASE_ 						=     share_att_key
											SCREAMING_SNAKE_CASE_ 						=     relative_attention
											SCREAMING_SNAKE_CASE_ 						=     norm_rel_ebd
											SCREAMING_SNAKE_CASE_ 						=     list(_lowercase		)
											SCREAMING_SNAKE_CASE_ 						=     hidden_act
											SCREAMING_SNAKE_CASE_ 						=     num_attention_heads
											SCREAMING_SNAKE_CASE_ 						=     hidden_dropout
											SCREAMING_SNAKE_CASE_ 						=     attention_dropout
											SCREAMING_SNAKE_CASE_ 						=     activation_dropout
											SCREAMING_SNAKE_CASE_ 						=     feat_proj_dropout
											SCREAMING_SNAKE_CASE_ 						=     final_dropout
											SCREAMING_SNAKE_CASE_ 						=     layer_norm_eps
											SCREAMING_SNAKE_CASE_ 						=     feature_layer_norm_eps
											SCREAMING_SNAKE_CASE_ 						=     initializer_range
											SCREAMING_SNAKE_CASE_ 						=     vocab_size

											if (
											    (len(self.conv_stride		) != self.num_feat_extract_layers)
											    or (len(self.conv_kernel		) != self.num_feat_extract_layers)
											    or (len(self.conv_dim		) != self.num_feat_extract_layers)
											):
																		raise ValueError(
																		    '''Configuration for convolutional layers is incorrect.'''
																		    '''It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`,'''
																		    f"but is `len(config.conv_dim) = {len(self.conv_dim		)}`, `len(config.conv_stride)"
																		    f"= {len(self.conv_stride		)}`, `len(config.conv_kernel) = {len(self.conv_kernel		)}`."		)

											# fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
											SCREAMING_SNAKE_CASE_ 						=     apply_spec_augment
											SCREAMING_SNAKE_CASE_ 						=     mask_time_prob
											SCREAMING_SNAKE_CASE_ 						=     mask_time_length
											SCREAMING_SNAKE_CASE_ 						=     mask_time_min_masks
											SCREAMING_SNAKE_CASE_ 						=     mask_feature_prob
											SCREAMING_SNAKE_CASE_ 						=     mask_feature_length
											SCREAMING_SNAKE_CASE_ 						=     mask_feature_min_masks

											# ctc loss
											SCREAMING_SNAKE_CASE_ 						=     ctc_loss_reduction
											SCREAMING_SNAKE_CASE_ 						=     ctc_zero_infinity

											# sequence classification
											SCREAMING_SNAKE_CASE_ 						=     use_weighted_layer_sum
											SCREAMING_SNAKE_CASE_ 						=     classifier_proj_size





				@property
				def      _A				(				self: Any		):
											return functools.reduce(operator.mul		,       self.conv_stride		,       1		)





import inspect
import logging
import os
import random
import shutil
import tempfile
import unittest

import pytest
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_cuda
from accelerate.utils import ProjectConfiguration, set_seed


SCREAMING_SNAKE_CASE		:Optional[int]						=     logging.getLogger(__name__)




def 			lowerCAmelCase(      SCREAMING_SNAKE_CASE_=2			,						SCREAMING_SNAKE_CASE_=3			,						SCREAMING_SNAKE_CASE_=1_6			,						SCREAMING_SNAKE_CASE_ = 1_0			,						SCREAMING_SNAKE_CASE_ = 2			)->						Optional[Any]:



    """simple docstring"""






    def get_dataset(SCREAMING_SNAKE_CASE_			):
        UpperCamelCase_			  =						torch.randn(batch_size * n_batches			,						1			)
        return TensorDataset(SCREAMING_SNAKE_CASE_			,						a * x + b + 0.1 * torch.randn(batch_size * n_batches			,						1			)			)

    UpperCamelCase_			  =						get_dataset(SCREAMING_SNAKE_CASE_			)
    UpperCamelCase_			  =						get_dataset(SCREAMING_SNAKE_CASE_			)
    UpperCamelCase_			  =						DataLoader(SCREAMING_SNAKE_CASE_			,						shuffle=SCREAMING_SNAKE_CASE_			,						batch_size=SCREAMING_SNAKE_CASE_			,						num_workers=4			)
    UpperCamelCase_			  =						DataLoader(SCREAMING_SNAKE_CASE_			,						shuffle=SCREAMING_SNAKE_CASE_			,						batch_size=SCREAMING_SNAKE_CASE_			,						num_workers=4			)
    return (train_dataloader, valid_dataloader)




def 			lowerCAmelCase(      SCREAMING_SNAKE_CASE_			,						SCREAMING_SNAKE_CASE_			,						SCREAMING_SNAKE_CASE_			,						SCREAMING_SNAKE_CASE_			,						SCREAMING_SNAKE_CASE_			,						SCREAMING_SNAKE_CASE_=None			)->						Any:



    """simple docstring"""






    UpperCamelCase_			  =						[]
    for epoch in range(SCREAMING_SNAKE_CASE_			):
        # Train quickly
        model.train()
        for batch in dataloader:
            UpperCamelCase_  ,       UpperCamelCase_			  =						batch
            UpperCamelCase_			  =						model(SCREAMING_SNAKE_CASE_			)
            UpperCamelCase_			  =						torch.nn.functional.mse_loss(SCREAMING_SNAKE_CASE_			,						SCREAMING_SNAKE_CASE_			)
            accelerator.backward(SCREAMING_SNAKE_CASE_			)
            optimizer.step()
            optimizer.zero_grad()
        rands.append(random.random()			)  # Introduce some randomness
        if scheduler is not None:
            scheduler.step()
    return rands






class 						__magic_name__					(							nn.Module				):






     def __init__(				self	)->     List[Any]:
         super().__init__()
         UpperCamelCase_			  =						nn.Parameter(torch.randn(1	)	)
         UpperCamelCase_			  =						nn.Parameter(torch.randn(1	)	)






     def 	UpperCAmelCase_				(				self		,						_lowercase	)->     str:
         return x * self.a + self.b






class 						__magic_name__					(							unittest.TestCase				):






     def 	UpperCAmelCase_				(				self	)->     List[Any]:
         with tempfile.TemporaryDirectory() as tmpdir:
             set_seed(42	)
             UpperCamelCase_			  =						DummyModel()
             UpperCamelCase_			  =						torch.optim.Adam(params=model.parameters()		,						lr=1e-3	)
             UpperCamelCase_  ,       UpperCamelCase_			  =						dummy_dataloaders()
             UpperCamelCase_			  =						ProjectConfiguration(total_limit=1		,						project_dir=_lowercase		,						automatic_checkpoint_naming=_lowercase	)
             # Train baseline
             UpperCamelCase_			  =						Accelerator(project_config=_lowercase	)
             UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_			  =						accelerator.prepare(
                 _lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             # Save initial
             accelerator.save_state()

             # Save second state
             accelerator.save_state()
             self.assertEqual(len(os.listdir(accelerator.project_dir	)	)		,						1	)






     def 	UpperCAmelCase_				(				self	)->     Optional[Any]:
         with tempfile.TemporaryDirectory() as tmpdir:
             set_seed(42	)
             UpperCamelCase_			  =						DummyModel()
             UpperCamelCase_			  =						torch.optim.Adam(params=model.parameters()		,						lr=1e-3	)
             UpperCamelCase_  ,       UpperCamelCase_			  =						dummy_dataloaders()
             # Train baseline
             UpperCamelCase_			  =						Accelerator()
             UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_			  =						accelerator.prepare(
                 _lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             # Save initial
             UpperCamelCase_			  =						os.path.join(_lowercase		,						"initial"	)
             accelerator.save_state(_lowercase	)
             ((UpperCamelCase_)  ,       (UpperCamelCase_))			  =						model.a.item(), model.b.item()
             UpperCamelCase_			  =						optimizer.state_dict()
             UpperCamelCase_			  =						train(3		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             ((UpperCamelCase_)  ,       (UpperCamelCase_))			  =						model.a.item(), model.b.item()
             UpperCamelCase_			  =						optimizer.state_dict()

             # Train partially
             set_seed(42	)
             UpperCamelCase_			  =						DummyModel()
             UpperCamelCase_			  =						torch.optim.Adam(params=model.parameters()		,						lr=1e-3	)
             UpperCamelCase_  ,       UpperCamelCase_			  =						dummy_dataloaders()
             UpperCamelCase_			  =						Accelerator()
             UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_			  =						accelerator.prepare(
                 _lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             accelerator.load_state(_lowercase	)
             ((UpperCamelCase_)  ,       (UpperCamelCase_))			  =						model.a.item(), model.b.item()
             UpperCamelCase_			  =						optimizer.state_dict()
             self.assertEqual(_lowercase		,						_lowercase	)
             self.assertEqual(_lowercase		,						_lowercase	)
             self.assertEqual(_lowercase		,						_lowercase	)

             UpperCamelCase_			  =						train(2		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             # Save everything
             UpperCamelCase_			  =						os.path.join(_lowercase		,						"checkpoint"	)
             accelerator.save_state(_lowercase	)

             # Load everything back in and make sure all states work
             accelerator.load_state(_lowercase	)
             test_rands += train(1		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             ((UpperCamelCase_)  ,       (UpperCamelCase_))			  =						model.a.item(), model.b.item()
             UpperCamelCase_			  =						optimizer.state_dict()
             self.assertEqual(_lowercase		,						_lowercase	)
             self.assertEqual(_lowercase		,						_lowercase	)
             self.assertEqual(_lowercase		,						_lowercase	)
             self.assertEqual(_lowercase		,						_lowercase	)






     def 	UpperCAmelCase_				(				self	)->     Union[str, Any]:
         with tempfile.TemporaryDirectory() as tmpdir:
             set_seed(42	)
             UpperCamelCase_			  =						DummyModel()
             UpperCamelCase_			  =						torch.optim.Adam(params=model.parameters()		,						lr=1e-3	)
             UpperCamelCase_  ,       UpperCamelCase_			  =						dummy_dataloaders()
             UpperCamelCase_			  =						ProjectConfiguration(automatic_checkpoint_naming=_lowercase	)

             # Train baseline
             UpperCamelCase_			  =						Accelerator(project_dir=_lowercase		,						project_config=_lowercase	)
             UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_			  =						accelerator.prepare(
                 _lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             # Save initial
             accelerator.save_state()
             ((UpperCamelCase_)  ,       (UpperCamelCase_))			  =						model.a.item(), model.b.item()
             UpperCamelCase_			  =						optimizer.state_dict()
             UpperCamelCase_			  =						train(3		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             ((UpperCamelCase_)  ,       (UpperCamelCase_))			  =						model.a.item(), model.b.item()
             UpperCamelCase_			  =						optimizer.state_dict()

             # Train partially
             set_seed(42	)
             UpperCamelCase_			  =						DummyModel()
             UpperCamelCase_			  =						torch.optim.Adam(params=model.parameters()		,						lr=1e-3	)
             UpperCamelCase_  ,       UpperCamelCase_			  =						dummy_dataloaders()
             UpperCamelCase_			  =						ProjectConfiguration(iteration=1		,						automatic_checkpoint_naming=_lowercase	)
             UpperCamelCase_			  =						Accelerator(project_dir=_lowercase		,						project_config=_lowercase	)
             UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_			  =						accelerator.prepare(
                 _lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             accelerator.load_state(os.path.join(_lowercase		,						"checkpoints"		,						"checkpoint_0"	)	)
             ((UpperCamelCase_)  ,       (UpperCamelCase_))			  =						model.a.item(), model.b.item()
             UpperCamelCase_			  =						optimizer.state_dict()
             self.assertEqual(_lowercase		,						_lowercase	)
             self.assertEqual(_lowercase		,						_lowercase	)
             self.assertEqual(_lowercase		,						_lowercase	)

             UpperCamelCase_			  =						train(2		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             # Save everything
             accelerator.save_state()

             # Load everything back in and make sure all states work
             accelerator.load_state(os.path.join(_lowercase		,						"checkpoints"		,						"checkpoint_1"	)	)
             test_rands += train(1		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             ((UpperCamelCase_)  ,       (UpperCamelCase_))			  =						model.a.item(), model.b.item()
             UpperCamelCase_			  =						optimizer.state_dict()
             self.assertEqual(_lowercase		,						_lowercase	)
             self.assertEqual(_lowercase		,						_lowercase	)
             self.assertEqual(_lowercase		,						_lowercase	)
             self.assertEqual(_lowercase		,						_lowercase	)






     def 	UpperCAmelCase_				(				self	)->     Any:
         UpperCamelCase_			  =						torch.tensor([1, 2, 3]	)
         UpperCamelCase_			  =						torch.tensor([2, 3, 4]	)
         UpperCamelCase_			  =						DummyModel()
         UpperCamelCase_			  =						torch.optim.Adam(net.parameters()	)
         UpperCamelCase_			  =						Accelerator()
         with self.assertRaises(_lowercase	) as ve:
             accelerator.register_for_checkpointing(_lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
         UpperCamelCase_			  =						str(ve.exception	)
         self.assertTrue("Item at index 0" in message	)
         self.assertTrue("Item at index 1" in message	)
         self.assertFalse("Item at index 2" in message	)
         self.assertFalse("Item at index 3" in message	)






     def 	UpperCAmelCase_				(				self	)->     Optional[Any]:
         with tempfile.TemporaryDirectory() as tmpdir:
             set_seed(42	)
             UpperCamelCase_			  =						DummyModel()
             UpperCamelCase_			  =						torch.optim.Adam(params=model.parameters()		,						lr=1e-3	)
             UpperCamelCase_			  =						torch.optim.lr_scheduler.StepLR(_lowercase		,						step_size=1		,						gamma=0.99	)
             UpperCamelCase_  ,       UpperCamelCase_			  =						dummy_dataloaders()
             UpperCamelCase_			  =						ProjectConfiguration(automatic_checkpoint_naming=_lowercase	)
             # Train baseline
             UpperCamelCase_			  =						Accelerator(project_dir=_lowercase		,						project_config=_lowercase	)
             UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_  ,       UpperCamelCase_			  =						accelerator.prepare(
                 _lowercase		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             # Save initial
             accelerator.save_state()
             UpperCamelCase_			  =						scheduler.state_dict()
             train(3		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase	)
             self.assertNotEqual(_lowercase		,						scheduler.state_dict()	)

             # Load everything back in and make sure all states work
             accelerator.load_state(os.path.join(_lowercase		,						"checkpoints"		,						"checkpoint_0"	)	)
             self.assertEqual(_lowercase		,						scheduler.state_dict()	)






     def 	UpperCAmelCase_				(				self	)->     str:
         with tempfile.TemporaryDirectory() as tmpdir:
             set_seed(42	)
             UpperCamelCase_			  =						DummyModel()
             UpperCamelCase_			  =						ProjectConfiguration(automatic_checkpoint_naming=_lowercase		,						total_limit=2	)
             # Train baseline
             UpperCamelCase_			  =						Accelerator(project_dir=_lowercase		,						project_config=_lowercase	)
             UpperCamelCase_			  =						accelerator.prepare(_lowercase	)
             # Save 3 states:
             for _ in range(11	):
                 accelerator.save_state()
             self.assertTrue(not os.path.exists(os.path.join(_lowercase		,						"checkpoints"		,						"checkpoint_0"	)	)	)
             self.assertTrue(os.path.exists(os.path.join(_lowercase		,						"checkpoints"		,						"checkpoint_9"	)	)	)
             self.assertTrue(os.path.exists(os.path.join(_lowercase		,						"checkpoints"		,						"checkpoint_10"	)	)	)






     @require_cuda
     def 	UpperCAmelCase_				(				self	)->     int:
         UpperCamelCase_			  =						["torchrun", F"--nproc_per_node={torch.cuda.device_count()}", inspect.getfile(self.__class__	)]
         execute_subprocess_async(_lowercase		,						env=os.environ.copy()	)


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE		:Any						=     """/tmp/accelerate/state_checkpointing"""
 SCREAMING_SNAKE_CASE		:Any						=     DummyModel()
 SCREAMING_SNAKE_CASE		:Union[str, Any]						=     torch.optim.Adam(params=model.parameters(), lr=1e-3)
 SCREAMING_SNAKE_CASE		:Dict						=     torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9_9)
 SCREAMING_SNAKE_CASE	,							SCREAMING_SNAKE_CASE		:str						=     dummy_dataloaders()
 SCREAMING_SNAKE_CASE		:Tuple						=     ProjectConfiguration(automatic_checkpoint_naming=True)
 # Train baseline
 SCREAMING_SNAKE_CASE		:int						=     Accelerator(project_dir=savedir, project_config=project_config, mixed_precision="""no""")
 if accelerator.process_index == 0:
  if os.path.exists(savedir):
   shutil.rmtree(savedir)
  os.makedirs(savedir)
 SCREAMING_SNAKE_CASE	,							SCREAMING_SNAKE_CASE	,							SCREAMING_SNAKE_CASE	,							SCREAMING_SNAKE_CASE	,							SCREAMING_SNAKE_CASE		:Tuple						=     accelerator.prepare(
     model, optimizer, train_dataloader, valid_dataloader, scheduler
 )
 SCREAMING_SNAKE_CASE	,							SCREAMING_SNAKE_CASE		:List[str]						=     accelerator.prepare(model, optimizer)
 train(3, model, train_dataloader, optimizer, accelerator, scheduler)
 # Check that the intial optimizer is loaded on the GPU
 for group in optimizer.param_groups:
  SCREAMING_SNAKE_CASE		:Optional[Any]						=     group["""params"""][0].device
  break
 assert param_device.type == accelerator.device.type
 SCREAMING_SNAKE_CASE		:List[str]						=     model.cpu()
 accelerator.wait_for_everyone()
 accelerator.save_state()
 accelerator.wait_for_everyone()

 # Check CPU state
 accelerator.load_state(os.path.join(savedir, """checkpoints""", """checkpoint_0"""), map_location="""cpu""")
 for group in optimizer.param_groups:
  SCREAMING_SNAKE_CASE		:Optional[int]						=     group["""params"""][0].device
  break
 assert (
     param_device.type == torch.device("""cpu""").type
 ), F"Loaded optimizer states did not match, expected to be loaded on the CPU but got {param_device}"

 # Check device state
 model.to(accelerator.device)
 accelerator.load_state(os.path.join(savedir, """checkpoints""", """checkpoint_0"""), map_location="""on_device""")
 for group in optimizer.param_groups:
  SCREAMING_SNAKE_CASE		:str						=     group["""params"""][0].device
  break
 assert (
     param_device.type == accelerator.device.type
 ), F"Loaded optimizer states did not match, expected to be loaded on {accelerator.device} but got {param_device}"

 # Check error
 with pytest.raises(TypeError, match="""Unsupported optimizer map location passed"""):
  accelerator.load_state(os.path.join(savedir, """checkpoints""", """checkpoint_0"""), map_location="""invalid""")
 accelerator.wait_for_everyone()
 if accelerator.process_index == 0:
  shutil.rmtree(savedir)
 accelerator.wait_for_everyone()
