code: def test_lr_schedule(example_experts):
      expert_backend = example_experts[EXPERT_NAME]
      optimizer = expert_backend.optimizer
  
      batch = torch.randn(1, 1)
      loss_grad = torch.randn(1, 1)
  
      with TemporaryDirectory() as tmpdir:
          tmp_path = Path(tmpdir)
  
          assert optimizer.param_groups[0]["lr"] == 0.0
  
          for i in range(BACKWARD_PASSES_BEFORE_SAVE):
              assert (
                  optimizer.param_groups[0]["lr"]
                  == PEAK_LR * i / BACKWARD_PASSES_BEFORE_SAVE
              )
              expert_backend.backward(batch, loss_grad)
  
          assert optimizer.param_groups[0]["lr"] == PEAK_LR
          store_experts(example_experts, tmp_path)
  
          for i in range(BACKWARD_PASSES_AFTER_SAVE):
              assert optimizer.param_groups[0]["lr"] == PEAK_LR * (
                  1 - (i / BACKWARD_PASSES_AFTER_SAVE)
              )
              expert_backend.backward(batch, loss_grad)
  
          assert optimizer.param_groups[0]["lr"] == 0.0
          load_experts(example_experts, tmp_path)
          assert optimizer.param_groups[0]["lr"] == PEAK_LR
  
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 4 for calling a global function with example usage: test_lr_schedule(example_experts) without return types