code: def test_multihead_expert(hid_dim=16):
      with background_server(
          expert_cls="multihead",
          num_experts=2,
          device="cpu",
          hidden_dim=hid_dim,
          num_handlers=2,
          custom_module_path=CUSTOM_EXPERTS_PATH,
      ) as server_peer_info:
          dht = DHT(initial_peers=server_peer_info.addrs, start=True)
          expert0, expert1 = create_remote_experts(
              [
                  ExpertInfo(
                      uid="expert.0", peer_id=server_peer_info.peer_id
                  ),
                  ExpertInfo(
                      uid="expert.1", peer_id=server_peer_info.peer_id
                  ),
              ],
              dht=dht,
          )
  
          for batch_size in (1, 4):
              batch = (
                  torch.randn(batch_size, hid_dim),
                  torch.randn(batch_size, 2 * hid_dim),
                  torch.randn(batch_size, 3 * hid_dim),
              )
  
              output0 = expert0(*batch)
              output1 = expert1(*batch)
  
              loss = output0.sum()
              loss.backward()
              loss = output1.sum()
              loss.backward()
  
file_path: hivemind/tests/test_custom_experts.py
imports: import os
  import pytest
  import torch
  from hivemind.dht import DHT
  from hivemind.moe.client.expert import create_remote_experts
  from hivemind.moe.expert_uid import ExpertInfo
  from hivemind.moe.server import background_server
label_desc: Use image node_id 2 for calling a global function with example usage: test_multihead_expert(hid_dim) without return types