import torch


def convert_half_precision(model):
    """
    Converts a torch model to half precision. Keeps the batch normalization layers
    at single precision
    source: https://github.com/onnx/onnx-tensorrt/issues/235#issuecomment-523948414
    """
    
    def bn_to_float(module):
        """
        BatchNorm layers need parameters in single precision. Find all layers and convert
        them back to float.
        """
        if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
            module.float()
        for child in module.children():
            bn_to_float(child)
        return module
    
    return bn_to_float(model.half())