def convert_dtype(dtype, obj):
    """Converts given tensor to given dtype

    Args:
        dtype (str): One of `fp32` or `fp64`
        obj (`obj`:torch.Tensor | `obj`:torch.nn.Module): Module or tensor to convert

    Returns:
        (`obj`:torch.Tensor | `obj`:torch.nn.Module): Converted tensor or module
    """
    # The object should be a ``module`` or a ``tensor``
    if dtype == "fp32":
        return obj.float()
    elif dtype == "fp64":
        return obj.double()
    else:
        raise NotImplementedError("dtype {} not supported.".format(dtype))