code: def do_test_forward(device):
      torch.manual_seed(3)
      input = torch.randn(12, 4).to(device)
      gate = Top2Gate(4, 6).to(device)
      capacity = 2 * 12 // 6
      l_aux, combine_weights, dispatch_mask = gate(input)
      assert pytest.approx(l_aux.item(), rel=0.01) == 0.0267, l_aux
      assert combine_weights.shape == (12, 6, 4)
      assert dispatch_mask.shape == (12, 6, 4)
      assert torch.equal(combine_weights.bool(), dispatch_mask)
      assert torch.all(
          torch.sum(dispatch_mask, axis=(0, 2)) <= capacity
      )
      assert torch.all(combine_weights >= 0.0)
      assert torch.all(combine_weights <= 1.0)
      weights_sum = torch.sum(combine_weights).item()
      assert round(weights_sum) == pytest.approx(
          weights_sum
      ), weights_sum
      # For this random seed, we get 12 slots filled.
      assert weights_sum == pytest.approx(12.0), weights_sum
  
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 3 for calling a global function with example usage: do_test_forward(device) without return types