def get_amp(fp16):
    """This function ensures that fp16 execution of torch.einsum is enabled
        if fp16 is set. Otherwise, it'll default to "promote" mode,
        where the operations are in fp32.
        Note that setting `fp16_opt_level="O2"` will remove the need for this code.
    """
    # Before we do anything with models, we want to
    if fp16:
        try:
            from apex import amp

            amp.register_half_function(torch, "einsum")
        except ImportError:
            raise ImportError(
                "Please install apex from https://www.github.com/nvidia/apex"
            )
    else:
        amp = None
    return amp
This function ensures that fp16 execution of torch.einsum is enabled
        if fp16 is set. Otherwise, it'll default to "promote" mode,
        where the operations are in fp32.
        Note that setting `fp16_opt_level="O2"` will remove the need for this code.
    
utils_nlp/common/pytorch_utils.py