code: def run(arguments):
      azure_info_path = arguments.get("--azure-info", None)
      training_data_path = RichPath.create(
          arguments["TRAIN_DATA_PATH"], azure_info_path
      )
      training_data = ShardedLazyDataIterable(
          partial(load_from_folder, training_data_path, shuffle=True)
      )
  
      validation_data_path = RichPath.create(
          arguments["VALID_DATA_PATH"], azure_info_path
      )
      validation_data = ShardedLazyDataIterable(
          partial(load_from_folder, validation_data_path, shuffle=False)
      )
  
      model_path = Path(arguments["MODEL_FILENAME"])
      assert model_path.name.endswith(
          ".pkl.gz"
      ), "MODEL_FILENAME must have a `.pkl.gz` suffix."
  
      initialize_metadata = True
      restore_path = arguments.get("--restore-path", None)
      if restore_path:
          initialize_metadata = False
          model, nn = Graph2Class.restore_model(
              Path(restore_path), device="cpu"
          )
      elif arguments["--autorestore"] and model_path.exists():
          initialize_metadata = False
          model, nn = Graph2Class.restore_model(
              model_path, device="cpu"
          )
      else:
          nn = None
          model = create_graph2class_gnn_model()
  
      if arguments["--restore-optimizer"]:
          opt_state = torch.load(arguments["--restore-optimizer"])
          current_epoch_idx = opt_state["epoch"]
          create_optimizer_ = partial(
              create_optimizer, state=opt_state["optimizer_state_dict"]
          )
      else:
          create_optimizer_ = create_optimizer
          current_epoch_idx = 0
  
      trainer = DistributedModelTrainer(
          model,
          model_path,
          max_num_epochs=int(arguments["--max-num-epochs"]),
          minibatch_size=int(arguments["--minibatch-size"]),
          optimizer_creator=create_optimizer_,
          clip_gradient_norm=1,
          target_validation_metric="Accuracy",
          target_validation_metric_higher_is_better=True,
          enable_amp=arguments["--amp"],
      )
      if nn is not None:
          trainer.neural_module = nn
  
      world_size = int(arguments["--world-size"])
      if world_size == -1:
          world_size = torch.cuda.device_count()
  
      trainer.distributed_train(
          world_size,
          training_data,
          validation_data,
          initialize_metadata=initialize_metadata,
          parallelize=not arguments["--sequential-run"],
          validate_on_start=True,
          shuffle_training_data=True,
          patience=10,
          worker_init=worker_init,
          start_epoch_idx=current_epoch_idx,
      )
  
imports: import random
  import torch
  import torch.distributed
  from docopt import docopt
  from dpu_utils.utils import RichPath, run_and_debug
  from functools import partial
  from pathlib import Path
  from torch.distributed.optim import ZeroRedundancyOptimizer
  from typing import Optional
  from ptgnn.baseneuralmodel.distributedtrainer import DistributedModelTrainer
  from ptgnn.baseneuralmodel.utils.amlutils import configure_logging, log_run
  from ptgnn.baseneuralmodel.utils.data import ShardedLazyDataIterable
  from ptgnn.implementations.typilus.graph2class import Graph2Class
  from ptgnn.implementations.typilus.train import create_graph2class_gnn_model
label_desc: Use image node_id 5 for calling a global function with example usage: run(arguments) without return types