def METHOD_NAME(self, model_context, config=None):
    if config is not None:
        half_dtype = config.get("dtype", None)
        if half_dtype in ("bf16", torch.bfloat16):
            if not torch.cuda.is_bf16_supported():
                gpu_model = torch.cuda.get_device_properties(torch.cuda.current_device()).name
                logger.info(f"Current device {gpu_model} does not support bfloat16. Set dtype to float16.")
                config["dtype"] = torch.float16
            else:
                config["dtype"] = torch.bfloat16
                config.setdefault("skip_if_nonfinite", False)
        elif half_dtype in ("fp16", torch.float16, None):
            config["dtype"] = torch.float16
        else:
            raise ValueError(
                "Invalid dtype. Supported dtypes are 'fp16', 'bf16', torch.float16, torch.bfloat16 but got"
                f" {half_dtype}"
            )
    else:
        config = {"enabled": True, "dtype": torch.float16}
    model_context.add_wrapper("amp_native", AmpNativeOptimization.apply_wrapper, config, is_pre_wrapper=False)
    return model_context