code: def test_non_pld_model(self):
      gamma = 0.001
      theta = 0.5
      config_dict = {
          "train_batch_size": 1,
          "steps_per_print": 1,
          "optimizer": {"type": "Adam", "params": {"lr": 0.0001}},
          "fp16": {"enabled": True},
          "progressive_layer_drop": {
              "enabled": True,
              "theta": theta,
              "gamma": gamma,
          },
      }
      hidden_dim = 10
  
      model = SimpleModel(hidden_dim, empty_grad=False)
      model, _, _, _ = deepspeed.initialize(
          config=config_dict,
          model=model,
          model_parameters=model.parameters(),
      )
  
      data_loader = random_dataloader(
          model=model,
          total_samples=1,
          hidden_dim=hidden_dim,
          device=model.device,
      )
  
      for i, batch in enumerate(data_loader):
          with pytest.raises(TypeError):
              loss = model(batch[0], batch[1])
  
file_path: H2O/h2o_flexgen/benchmark/third_party/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 the TestNonPLDModel obj's underlying member method code with example usage: obj.test_non_pld_model() without return types