code: def dynamic_partition(
      data: Tensor, partitions: Tensor, num_partitions: int
  ) -> List[Tensor]:
      """Implements tf.dynamic_partition in torch.
  
      From https://discuss.pytorch.org/t/equivalent-of-tf-dynamic-partition/53735.
      """
      assert data.size() == partitions.size()
  
      # Flatten data into 1D vectors to do partitioning correctly.
      data = data.view(-1)
      partitions = partitions.view(-1)
      result = []
      for i in range(num_partitions):
          result += [data[(partitions == i).nonzero().squeeze(1)]]
      return result
  
file_path: ludwig/ludwig/utils/metric_utils.py
imports: from collections import defaultdict, namedtuple
  from typing import Dict, List, Optional
  import torch
  from torch import Tensor
  from torchmetrics.metric import Metric
  from ludwig.constants import COMBINED, LOSS, NAME, TYPE
  from ludwig.modules.metric_registry import get_metric_names_for_type
  from ludwig.types import FeatureConfigDict
label_desc: Use image node_id 2 for calling a global function with example usage: dynamic_partition(data, partitions, num_partitions) and returns: result