code: def test_restore_update_count(example_experts):
      expert_backend = example_experts[EXPERT_NAME]
  
      batch = torch.randn(1, 1)
      loss_grad = torch.randn(1, 1)
  
      with TemporaryDirectory() as tmpdir:
          tmp_path = Path(tmpdir)
  
          for _ in range(BACKWARD_PASSES_BEFORE_SAVE):
              expert_backend.backward(batch, loss_grad)
  
          store_experts(example_experts, tmp_path)
  
          for _ in range(BACKWARD_PASSES_AFTER_SAVE):
              expert_backend.backward(batch, loss_grad)
  
          load_experts(example_experts, tmp_path)
          assert (
              expert_backend.scheduler._step_count
              == BACKWARD_PASSES_BEFORE_SAVE + 1
          )
  
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 3 for calling a global function with example usage: test_restore_update_count(example_experts) without return types