class_docstr: A subclass to aid in serialization. This allows us to save only the name of the
  dataset and reload it from TorchVision at deserialization time.
class_docstr_tok: ["A","subclass","to","aid","in","serialization",".","This","allows","us","to","save","only","the","name","of","the","dataset","and","reload","it","from","TorchVision","at","deserialization","time","."]
code: def deserialize_init_args(
      cls,
      args: Dict[str, Any],
      decoder_registry: Optional[TDecoderRegistry] = None,
      class_decoder_registry: Optional[TClassDecoderRegistry] = None,
  ) -> Dict[str, Any]:
      name = args["name"]
  
      dataset_fn = _REGISTRY[name]
  
      train_set = dataset_fn(
          root="./data",
          train=True,
          download=True,
          transform=transforms.ToTensor(),
      )
  
      test_set = dataset_fn(
          root="./data",
          train=False,
          download=True,
          transform=transforms.ToTensor(),
      )
  
      return {
          "name": name,
          "train_set": train_set,
          "test_set": test_set,
      }
  
file_path: Ax/ax/benchmark/problems/hpo/torchvision.py
imports: import os
  from typing import Any, Dict, Optional
  from ax.benchmark.problems.hpo.pytorch_cnn import PyTorchCNNBenchmarkProblem, PyTorchCNNRunner
  from ax.exceptions.core import UserInputError
  from ax.utils.common.serialization import TClassDecoderRegistry, TDecoderRegistry
  from ax.utils.common.typeutils import checked_cast
  from torch.utils.data import TensorDataset
label_desc: Use image node_id 2 for calling the PyTorchCNNTorchvisionRunner obj's underlying member method code with example usage: obj.deserialize_init_args(cls, args, decoder_registry, class_decoder_registry) and returns: dict