def set_amp_args(self, amp_args: Optional[Dict[(str, Any)]]):
    'Disable / enable apex.amp and set the automatic mixed precision parameters.\n\n        apex.amp can be utilized for mixed / half precision training.\n\n        Args:\n            amp_args: Dictionary containing arguments to be passed to\n            amp.initialize. Set to None to disable amp.  To enable mixed\n            precision training, pass amp_args={"opt_level": "O1"} here.\n            See https://nvidia.github.io/apex/amp.html for more info.\n\n        Raises:\n            RuntimeError: If opt_level is not None and apex is not installed.\n\n        Warning: apex needs to be installed to utilize this feature.\n        '
    self.amp_args = amp_args
    if (amp_args is None):
        logging.info('AMP disabled')
    else:
        try:
            self.amp_type = AmpType[self.amp_args['amp_type'].upper()]
        except KeyError:
            logging.info('AMP type not specified, defaulting to Apex')
            self.amp_type = AmpType.APEX
        if (not torch.cuda.is_available()):
            raise RuntimeError('AMP is required but CUDA is not supported, cannot enable AMP')
        if ((self.amp_type == AmpType.APEX) and (not apex_available)):
            raise RuntimeError('Apex AMP is required but Apex is not installed, cannot enable AMP')
        if self.use_sharded_ddp:
            if (self.amp_type == AmpType.APEX):
                raise RuntimeError('ShardedDDP has been requested, which is incompatible with Apex AMP')
            if (not fairscale_available):
                raise RuntimeError('ShardedDDP has been requested, but fairscale is not installed in the current environment')
        elif (self.amp_type == AmpType.PYTORCH):
            if self.use_sharded_ddp:
                logging.info('Using ShardedGradScaler to manage Pytorch AMP')
                self.amp_grad_scaler = ShardedGradScaler()
            else:
                self.amp_grad_scaler = TorchGradScaler()
        logging.info(f'AMP enabled with args {amp_args}')
    return self
Disable / enable apex.amp and set the automatic mixed precision parameters.

apex.amp can be utilized for mixed / half precision training.

Args:
    amp_args: Dictionary containing arguments to be passed to
    amp.initialize. Set to None to disable amp.  To enable mixed
    precision training, pass amp_args={"opt_level": "O1"} here.
    See https://nvidia.github.io/apex/amp.html for more info.

Raises:
    RuntimeError: If opt_level is not None and apex is not installed.

Warning: apex needs to be installed to utilize this feature.
def set_amp_args(self, amp_args: Optional[Dict[(str, Any)]]):
    'Disable / enable apex.amp and set the automatic mixed precision parameters.\n\n        apex.amp can be utilized for mixed / half precision training.\n\n        Args:\n            amp_args: Dictionary containing arguments to be passed to\n            amp.initialize. Set to None to disable amp.  To enable mixed\n            precision training, pass amp_args={"opt_level": "O1"} here.\n            See https://nvidia.github.io/apex/amp.html for more info.\n\n        Raises:\n            RuntimeError: If opt_level is not None and apex is not installed.\n\n        Warning: apex needs to be installed to utilize this feature.\n        '
    self.amp_args = amp_args
    if (amp_args is None):
        logging.info('AMP disabled')
    else:
        try:
            self.amp_type = AmpType[self.amp_args['amp_type'].upper()]
        except KeyError:
            logging.info('AMP type not specified, defaulting to Apex')
            self.amp_type = AmpType.APEX
        if (not torch.cuda.is_available()):
            raise RuntimeError('AMP is required but CUDA is not supported, cannot enable AMP')
        if ((self.amp_type == AmpType.APEX) and (not apex_available)):
            raise RuntimeError('Apex AMP is required but Apex is not installed, cannot enable AMP')
        if self.use_sharded_ddp:
            if (self.amp_type == AmpType.APEX):
                raise RuntimeError('ShardedDDP has been requested, which is incompatible with Apex AMP')
            if (not fairscale_available):
                raise RuntimeError('ShardedDDP has been requested, but fairscale is not installed in the current environment')
        elif (self.amp_type == AmpType.PYTORCH):
            if self.use_sharded_ddp:
                logging.info('Using ShardedGradScaler to manage Pytorch AMP')
                self.amp_grad_scaler = ShardedGradScaler()
            else:
                self.amp_grad_scaler = TorchGradScaler()
        logging.info(f'AMP enabled with args {amp_args}')
    return self