from typing import Union
import torch


def int2precision(precision: Union[int, torch.dtype]):
    """
    Get torch floating point precision from integer.
    If an instance of torch.dtype is passed, it is returned automatically.

    Args:
        precision (int, torch.dtype): Target precision.

    Returns:
        torch.dtupe: Floating point precision.
    """
    if isinstance(precision, torch.dtype):
        return precision
    else:
        try:
            return getattr(torch, f"float{precision}")
        except AttributeError:
            raise AttributeError(f"Unknown float precision {precision}")