class_docstr: Allows segregated segmentation in case the category of an object is known. This is the case in ShapeNet
  for example.
  
      Arguments:
          in_features -- size of the input channel
          cat_to_seg {[type]} -- category to segment maps for example:
              {
                  'Airplane': [0,1,2],
                  'Table': [3,4]
              }
  
      Keyword Arguments:
          dropout_proba  (default: {0.5})
          bn_momentum  -- batch norm momentum (default: {0.1})
      
code: def forward(self, features, category_labels, **kwargs):
      assert features.dim() == 2
      self._shifts = self._shifts.to(features.device)
      in_dim = features.shape[-1]
      features = self.channel_rasing(features)
      features = features.reshape((-1, self._num_categories, in_dim))
      features = features.transpose(
          0, 1
      )  # [num_categories, num_points, in_dim]
      features = (
          self.classifier(features) + self._bias
      )  # [num_categories, num_points, max_seg]
      ind = (
          category_labels.unsqueeze(-1)
          .repeat(1, 1, features.shape[-1])
          .long()
      )
  
      logits = features.gather(0, ind).squeeze(0)
      softmax = torch.nn.functional.log_softmax(logits, dim=-1)
  
      output = torch.zeros(logits.shape[0], self._max_seg + 1).to(
          features.device
      )
      cats_in_batch = torch.unique(category_labels)
      for cat in cats_in_batch:
          cat_mask = category_labels == cat
          seg_indices = self._cat_to_seg[cat.item()]
          probs = softmax[cat_mask, : len(seg_indices)]
          output[cat_mask, seg_indices[0] : seg_indices[-1] + 1] = probs
      return output
  
imports: import numpy
  import torch
  from torch import nn
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 2 for calling the MultiHeadClassifier obj's underlying member method code with example usage: obj.forward(features, category_labels) and returns: output