class_docstr: Segment Sum MLP model.
  
  Parameters
  ----------
  input_dim : int
      The input dim for the model.
  hidden_dim : int
      The hidden dim for the model.
  output_dim : int
      The output dim for the model.
  use_norm : bool
      Whether to normalize the segment sum or not.
  use_sigmoid : bool
      Whether to use sigmoid on the final output or not.
class_docstr_tok: ["Segment","Sum","MLP","model",".","Parameters","--","--","--","--","--","input_dim",":","int","The","input","dim","for","the","model",".","hidden_dim",":","int","The","hidden","dim","for","the","model",".","output_dim",":","int","The","output","dim","for","the","model",".","use_norm",":","bool","Whether","to","normalize","the","segment","sum","or","not",".","use_sigmoid",":","bool","Whether","to","use","sigmoid","on","the","final","output","or","not","."]
code: def forward(  # pylint: disable=missing-function-docstring
      self,
      segment_sizes: "torch.Tensor",
      features: "torch.Tensor",
  ) -> "torch.Tensor":
      n_seg = len(segment_sizes)
      encoded_features = self.encoder(features)
      segment_indices = torch.repeat_interleave(
          torch.arange(n_seg, device=features.device),
          segment_sizes.long(),
      )
      n_dim = encoded_features.shape[1]
      segment_sum = torch.scatter_add(
          input=torch.zeros(
              (n_seg, n_dim),
              dtype=encoded_features.dtype,
              device=features.device,
          ),
          dim=0,
          index=segment_indices.view(-1, 1).expand(-1, n_dim),
          src=encoded_features,
      )
      out = self.norm(segment_sum)
      out = self.layer0(out) + out
      out = self.layer1(out) + out
      out = self.decoder(out).squeeze()
      out = self.sigmoid(out)
      return out
  
imports: import glob
  import math
  import os
  import random
  import tempfile
  from collections import OrderedDict
  from itertools import chain
  from typing import Dict, List, NamedTuple, Optional, Tuple
  import numpy
  import torch
  import tvm
  from ...contrib.tar import tar, untar
  from ...runtime import NDArray
  from ...target import Target
  from ..cost_model import PyCostModel
  from ..database import JSONDatabase
  from ..feature_extractor import FeatureExtractor, PerStoreFeature
  from ..logging import get_logger
  from ..runner import RunnerResult
  from ..search_strategy import MeasureCandidate
  from ..tune_context import TuneContext
  from ..utils import derived_object, shash2hex
label_desc: Use image node_id 2 for calling the SegmentSumMLP obj's underlying member method code with example usage: obj.forward(segment_sizes, features) and returns: out