code: def _dython_theils_u(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)
  
file_path: torchmetrics/tests/unittests/nominal/test_theils_u.py
imports: import itertools
  import operator
  from functools import partial
  import pytest
  import torch
  from dython.nominal import theils_u
  from lightning_utilities.core.imports import compare_version
  from torchmetrics.functional.nominal.theils_u import theils_u, theils_u_matrix
  from torchmetrics.nominal import TheilsU
  from unittests import BATCH_SIZE, NUM_BATCHES, _Input
  from unittests.helpers.testers import MetricTester
label_desc: Use image node_id 2 for calling a global function with example usage: _dython_theils_u(preds, target, nan_strategy, nan_replace_value) and returns: torch