code: def test_expert1_overflow():
      num_tokens = 8
      num_experts = 4
      logits = torch.randn(num_tokens, num_experts)
      logits[:, 0] = (
          torch.max(logits, dim=1).values + 1
      )  # Force overflow
      top1s = torch.argmax(logits, dim=1)
      assert top1s.eq(0).all(), top1s
      _, __, dispatch_mask = top2gating(logits)
      capacity = 2 * num_tokens // num_experts
  
      for i in range(num_tokens):
          if i < capacity:
              assert dispatch_mask[i][0][i]
          else:
              assert not dispatch_mask[i][0].any()
  
file_path: fairscale/tests/nn/moe/test_top2gating.py
imports: import pytest
  import torch
  from fairscale.nn import Top2Gate
  from fairscale.nn.moe.top2gate import top2gating
label_desc: Use image node_id 6 for calling a global function with example usage: test_expert1_overflow() without return types