code: def test_moe(gate, number_of_local_experts, expert_constructor):
      test_config = {
          "name": "MixtureOfExperts",
          "dim_model": LATENT,
          "dropout": DROPOUT,
          "activation": Activation.ReLU,
          "hidden_layer_multiplier": 4,
          "number_of_experts": 4,
          "number_of_local_experts": number_of_local_experts,
          "gate": gate,
          "expert_constructor": expert_constructor,
      }
  
      init_torch_distributed_local()
  
      # dummy, just check construction and dimensions in the FW pass
      ffw = build_feedforward(test_config)
  
      inputs = torch.rand(
          BATCH, SEQ, LATENT, device=torch.device("cuda")
      )
      ffw = ffw.to(torch.device("cuda"))
  
      outputs = ffw(inputs)
      loss = torch.sum(outputs)
      loss.backward()
  
file_path: xformers/tests/test_feedforward.py
imports: import pytest
  import torch
  from xformers.components import Activation
  from xformers.components.feedforward import FEEDFORWARD_REGISTRY, build_feedforward
  from xformers.components.feedforward.mixture_of_experts import GateConfig
  from xformers.helpers.test_utils import init_torch_distributed_local
label_desc: Use image node_id 3 for calling a global function with example usage: test_moe(gate, number_of_local_experts, expert_constructor) without return types