def METHOD_NAME(preds, target, nan_strategy, nan_replace_value):
    preds = preds.argmax(1) if preds.ndim == 2 else preds
    target = target.argmax(1) if target.ndim == 2 else target
    v = dython_theils_u(
        preds.numpy(),
        target.numpy(),
        nan_strategy=nan_strategy,
        nan_replace_value=nan_replace_value,
    )
    return torch.tensor(v)