code: def test_training(max_steps: int = 100, threshold: float = 0.9):
      dataset = load_digits(n_class=2)
      X_train, y_train = torch.tensor(
          dataset["data"], dtype=torch.float
      ), torch.tensor(dataset["target"])
      SGD = partial(torch.optim.SGD, lr=0.05)
  
      with background_server(
          num_experts=2,
          device="cpu",
          optim_cls=SGD,
          hidden_dim=64,
          num_handlers=1,
          clip_grad_norm=1.0,
      ) as server_peer_info:
          dht = DHT(initial_peers=server_peer_info.addrs, start=True)
          expert1, expert2 = 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,
          )
          model = nn.Sequential(
              expert2, nn.ReLU(), expert1, nn.Linear(64, 2)
          )
  
          opt = SGD(model.parameters(), lr=0.05)
  
          for step in range(max_steps):
              outputs = model(X_train)
              loss = F.cross_entropy(outputs, y_train)
              loss.backward()
              opt.step()
              opt.zero_grad()
  
              accuracy = (
                  (outputs.argmax(dim=1) == y_train)
                  .float()
                  .mean()
                  .item()
              )
              if accuracy >= threshold:
                  break
  
          assert (
              accuracy >= threshold
          ), f"too small accuracy: {accuracy}"
  
file_path: hivemind/tests/test_training.py
imports: from functools import partial
  import pytest
  import torch
  import torch.nn
  import torch.nn.functional
  from sklearn.datasets import load_digits
  from hivemind import DHT
  from hivemind.moe.client import RemoteMixtureOfExperts, RemoteSwitchMixtureOfExperts
  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 1 for calling a global function with example usage: test_training(max_steps, threshold) without return types