code: def gshard_loss(scores_w_noise, top_ids):
      num_samples, num_global_experts = int(
          scores_w_noise.size(0)
      ), int(scores_w_noise.size(1))
      mask = _one_hot_with_dtype(
          top_ids[:, 0],
          num_global_experts,
          dtype=scores_w_noise.dtype,
          hot_value=num_global_experts / num_samples,
      )
      me = torch.sum(scores_w_noise, dim=0)
      ce = torch.sum(mask, dim=0)
      l_aux = torch.sum(me * ce) / num_samples
      return l_aux
  
imports: import torch
  from torch.distributions.normal import Normal
label_desc: Use image node_id 2 for calling a global function with example usage: gshard_loss(scores_w_noise, top_ids) and returns: l_aux