code: def test_save_load_checkpoints(example_experts):
      expert = example_experts[EXPERT_NAME].module
  
      with TemporaryDirectory() as tmpdir:
          tmp_path = Path(tmpdir)
  
          for i in range(1, EXPERT_WEIGHT_UPDATES + 1):
              expert.weight.data[0] = i
              store_experts(example_experts, tmp_path)
  
          checkpoints_dir = tmp_path / EXPERT_NAME
  
          assert checkpoints_dir.exists()
          # include checkpoint_last.pt
          assert (
              len(list(checkpoints_dir.iterdir()))
              == EXPERT_WEIGHT_UPDATES + 1
          )
  
          expert.weight.data[0] = 0
  
          load_experts(example_experts, tmp_path)
          assert expert.weight.data[0] == EXPERT_WEIGHT_UPDATES
  
file_path: hivemind/tests/test_expert_backend.py
imports: from pathlib import Path
  from tempfile import TemporaryDirectory
  import pytest
  import torch
  from torch.nn import Linear
  from hivemind import BatchTensorDescriptor, ModuleBackend
  from hivemind.moe.server.checkpoints import load_experts, store_experts
  from hivemind.moe.server.layers.lr_schedule import get_linear_schedule_with_warmup
label_desc: Use image node_id 2 for calling a global function with example usage: test_save_load_checkpoints(example_experts) without return types