code: def test_pld_schedule(tmpdir, theta):
      gamma = 0.001
  
      pld_scheduler = ProgressiveLayerDrop(theta, gamma)
      for i in range(10):
          pld_scheduler.update_state(i)
          expected_theta = (1.0 - theta) * np.exp(-gamma * i) + theta
          actual_theta = pld_scheduler.get_theta()
          assert expected_theta == actual_theta
  
file_path: DeepSpeed/tests/unit/runtime/test_pld.py
imports: import numpy
  import deepspeed
  import pytest
  from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop
  from unit.common import DistributedTest
  from unit.simple_model import SimpleModel, PLD_SimpleModel, random_dataloader
label_desc: Use image node_id 1 for calling a global function with example usage: test_pld_schedule(tmpdir, theta) without return types