code: def forward_multi(device):
      make_cudnn_deterministic()
      torch.set_printoptions(threshold=5000)
      num_local_experts = 4
      model_dim = 4
      num_experts = (
          dist.get_world_size(dist.group.WORLD) * num_local_experts
      )
      input = torch.randn(4 * num_local_experts, 16, model_dim).to(
          device
      )
      gate = Top2Gate(model_dim, num_experts)
      experts = []
      for i in range(num_local_experts):
          expert = torch.nn.Linear(model_dim, model_dim, bias=False)
          # Use identity matrix
          expert.weight = torch.nn.Parameter(torch.eye(model_dim))
          experts += [expert]
      moe = MOELayer(gate, torch.nn.ModuleList(experts)).to(device)
      output = moe(input)
      assert (
          output.shape == input.shape
      ), f"{output.shape} != {input.shape}"
      # 90% of the input should have gone to an expert
      assert (
          len(output.nonzero(as_tuple=False)) / output.numel() > 0.90
      ), f"{len(output.nonzero(as_tuple=False))} / {output.numel()}"
      # Except for zeros, re-assembled output should match input due to identity expert.
      torch.testing.assert_allclose(
          input, torch.where(output > 0, output, input)
      )
  
imports: import functools
  import tempfile
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  from fairscale.fair_dev.testing.testing import make_cudnn_deterministic
  from fairscale.internal import torch_version
  from fairscale.nn import MOELayer, Top2Gate
label_desc: Use image node_id 6 for calling a global function with example usage: forward_multi(device) without return types