code: def forward(self, input_features):
      features = input_features.reshape(-1, input_features.size(-1))
      is_training = input_features.requires_grad
  
      # if self.shuffle and is_training:
      #     # Send each token to a random worker, to break correlations within the batch
      #     shuffle_sort = torch.randperm(features.size(0), device=features.device)
      #     features = All2All.apply(features[shuffle_sort])
  
      tis, tvs = [], []
      for i in range(self.topk):
          print(i)
          with torch.no_grad():
              # Compute similarity of each token to each expert, for routing
              token_expert_affinities = features.matmul(
                  self.expert_centroids[i].transpose(0, 1)
              )
          print(token_expert_affinities)
  
          # Compute which token goes to which expert
          top_idx = self.balanced_assignment(token_expert_affinities)
          top_idx = np.delete(top_idx.cpu(), -1, axis=1).cuda()
  
          print(top_idx)
  
          # Swap these tokens for the right ones for our expert
          top_value = torch.empty(top_idx.size())
          for i in range(0, self.num_expert):
              for j in range(
                  0, int(int(features.size(0)) / self.num_expert)
              ):
                  top_value[i][j] = token_expert_affinities[
                      top_idx[i][j]
                  ][i]
          tis.append(top_idx)
          tvs.append(top_value)
  
      return tis, tvs
  
file_path: thu-pacman-faster-moe/fmoe/gates/base_layer_gate.py
imports: import torch.nn
  import torch
  import fmoe_cuda
  import numpy
label_desc: Use image node_id 2 for calling the BaseLayerGate obj's underlying member method code with example usage: obj.forward(input_features) and returns: tis,  tvs