code: def load_checkpoint(self, checkpoint: Union[Dict, str]):
      if self.return_type == "object":
          assert isinstance(checkpoint, dict)
          checkpoint_data = checkpoint
          checkpoint_file = None
      elif self.return_type == "root":
          assert "subdir" not in checkpoint
          checkpoint_file = os.path.join(
              checkpoint, "subdir", "checkpoint.pkl"
          )
      elif self.return_type == "subdir":
          assert "subdir" in checkpoint
          assert "checkpoint.pkl" not in checkpoint
          checkpoint_file = os.path.join(checkpoint, "checkpoint.pkl")
      else:  # self.return_type == "checkpoint"
          assert checkpoint.endswith("subdir/checkpoint.pkl")
          checkpoint_file = checkpoint
  
      if checkpoint_file:
          with open(checkpoint_file, "rb") as f:
              checkpoint_data = json.load(f)
  
      checkpoint_data = {
          key: value
          for key, value in checkpoint_data.items()
          if not key.startswith("_")
      }
      assert checkpoint_data == {"data": 1}, checkpoint_data
  
imports: import json
  import os
  from typing import Dict, Union
  import pytest
  import ray
  from ray import train, tune
  from ray.train._internal.storage import StorageContext
  from ray.tune.trainable import wrap_function
  from ray.train.tests.util import create_dict_checkpoint
label_desc: Use image node_id 4 for calling the SavingTrainable obj's underlying member method code with example usage: obj.load_checkpoint(checkpoint) without return types