code: def test_save_and_restore(self):
      estimator1 = Estimator.from_torch(
          model=model_creator,
          optimizer=optimizer_creator,
          loss=nn.MSELoss(),
          scheduler_creator=scheduler_creator,
          config={
              "lr": 1e-2,  # used in optimizer_creator
              "hidden_size": 1,  # used in model_creator
          },
          backend="horovod",
      )
      with TemporaryDirectory() as tmp_path:
          estimator1.fit(train_data_creator, batch_size=4, epochs=1)
          checkpoint_path = os.path.join(tmp_path, "checkpoint")
          estimator1.save(checkpoint_path)
  
          model1 = estimator1.get_model()
  
          estimator1.shutdown()
  
          estimator2 = Estimator.from_torch(
              model=model_creator,
              optimizer=optimizer_creator,
              loss=nn.MSELoss(),
              scheduler_creator=scheduler_creator,
              config={
                  "lr": 1e-2,  # used in optimizer_creator
                  "hidden_size": 1,  # used in model_creator
              },
              backend="horovod",
          )
          estimator2.load(checkpoint_path)
  
          model2 = estimator2.get_model()
  
      model1_state_dict = model1.state_dict()
      model2_state_dict = model2.state_dict()
  
      assert set(model1_state_dict.keys()) == set(
          model2_state_dict.keys()
      )
  
      for k in model1_state_dict:
          assert torch.equal(model1_state_dict[k], model2_state_dict[k])
      estimator2.shutdown()
  
file_path: analytics-zoo/pyzoo/test/zoo/orca/learn/ray/pytorch/test_estimator_horovod_backend.py
imports: from zoo.examples.orca.learn.horovod.pytorch_estimator import model_creator, optimizer_creator, scheduler_creator, train_data_creator, validation_data_creator
  from zoo.orca.learn.pytorch import Estimator
  from unittest import TestCase
  import pytest
  import torch
  import torch.nn
  import os
  from tempfile import TemporaryDirectory
label_desc: Use image node_id 3 for calling the TestPyTorchEstimator obj's underlying member method code with example usage: obj.test_save_and_restore() without return types