code: def update(
      self,
      indices_,
      locations_,
      gates_,
      capacity=None,
      is_postscore=True,
  ):
      self.indices_ = [x.to(torch.int32).view(-1) for x in indices_]
      self.locations_ = [x.to(torch.int32) for x in locations_]
      self.gates_ = [x.to(self.dtype) for x in gates_]
      self.is_postscore = is_postscore
      self.sample_size, self.capacity = (
          int(self.indices_[0].size(0)),
          int(capacity) or self.capacity,
      )
  
      if self.is_cuda != indices_[0].is_cuda:
          self.is_cuda = indices_[0].is_cuda
          if self.is_cuda not in TutelMoeFastDispatcher.kernel_pool:
              self.func_fwd = jit_kernel.create_forward(
                  self.dtype, indices_[0].is_cuda
              )
              self.func_bwd_data = jit_kernel.create_backward_data(
                  self.dtype, indices_[0].is_cuda
              )
              self.func_bwd_gate = jit_kernel.create_backward_gate(
                  self.dtype, indices_[0].is_cuda
              )
              TutelMoeFastDispatcher.kernel_pool[self.is_cuda] = (
                  self.func_fwd,
                  self.func_bwd_data,
                  self.func_bwd_gate,
              )
          else:
              (
                  self.func_fwd,
                  self.func_bwd_data,
                  self.func_bwd_gate,
              ) = TutelMoeFastDispatcher.kernel_pool[self.is_cuda]
  
      if (
          TutelMoeFastDispatcher.ones_helper is None
          or TutelMoeFastDispatcher.ones_helper.size(0)
          < self.sample_size
      ):
          TutelMoeFastDispatcher.ones_helper = torch.ones(
              [self.sample_size, 2],
              dtype=self.dtype,
              device=self.indices_[0].device,
          )
      if (
          TutelMoeFastDispatcher.ones_helper.is_cuda
          != self.indices_[0].is_cuda
      ):
          TutelMoeFastDispatcher.ones_helper = torch.ones(
              [TutelMoeFastDispatcher.ones_helper.size(0), 2],
              dtype=self.dtype,
              device=self.indices_[0].device,
          )
      self.ones_helper = TutelMoeFastDispatcher.ones_helper
  
imports: from typing import TYPE_CHECKING, Any, Optional, Tuple, Union, cast
  import logging
  import torch
  from torch import Tensor
  from .jit_compiler import IS_HIP_EXTENSION
  from ..jit_kernels import sparse
  from ..jit_kernels.gating import fast_cumsum_sub_one
  from .communicate import get_world_rank, simple_all_reduce
  from .None import losses
label_desc: Use image node_id 2 for calling the TutelMoeFastDispatcher obj's underlying member method code with example usage: obj.update(indices_, locations_, gates_, capacity, is_postscore) without return types