code: def run(arguments):
      if arguments["--aml"]:
          from azureml.core.run import Run
  
          aml_ctx = Run.get_context()
          assert (
              torch.cuda.is_available()
          ), "No CUDA available. Aborting training."
      else:
          aml_ctx = None
  
      log_path = configure_logging(aml_ctx)
      azure_info_path = arguments.get("--azure-info", None)
      training_data_path = RichPath.create(
          arguments["TRAIN_DATA_PATH"], azure_info_path
      )
      training_data = LazyDataIterable(
          lambda: load_from_folder(training_data_path, shuffle=True)
      )
  
      validation_data_path = RichPath.create(
          arguments["VALID_DATA_PATH"], azure_info_path
      )
      validation_data = LazyDataIterable(
          lambda: 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["--aml"] 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"]
  
          def create_optimizer(parameters):
              opt = torch.optim.Adam(parameters, lr=0.00025)
              logging.info(
                  "Restoring optimizer state from `%s`.",
                  arguments["--restore-optimizer"],
              )
  
              opt.load_state_dict(opt_state["optimizer_state_dict"])
              return opt
  
      else:
          current_epoch_idx = 0
  
          def create_optimizer(parameters):
              return torch.optim.Adam(parameters, lr=0.00025)
  
      trainer = ModelTrainer(
          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
  
      trainer.register_train_epoch_end_hook(
          lambda model, nn, epoch, metrics: log_run(
              aml_ctx, "train", model, epoch, metrics
          )
      )
      trainer.register_validation_epoch_end_hook(
          lambda model, nn, epoch, metrics: log_run(
              aml_ctx, "valid", model, epoch, metrics
          )
      )
  
      trainer.train(
          training_data,
          validation_data,
          show_progress_bar=not arguments["--quiet"],
          initialize_metadata=initialize_metadata,
          parallelize=not arguments["--sequential-run"],
          patience=10,
          store_tensorized_data_in_memory=True,
          start_epoch_idx=current_epoch_idx,
      )
  
      test_data_path = RichPath.create(
          arguments["TEST_DATA_PATH"], azure_info_path
      )
      test_data = LazyDataIterable(
          lambda: load_from_folder(test_data_path, shuffle=False)
      )
      acc = model.report_accuracy(
          test_data,
          trainer.neural_module,
          device=torch.device(
              "cuda:0" if torch.cuda.is_available() else "cpu"
          ),
      )
      print(f"Test accuracy: {acc:%}")
  
      if aml_ctx is not None:
          aml_ctx.log("Test Accuracy", acc)
          aml_ctx.upload_file(
              name="model.pkl.gz", path_or_stream=str(model_path)
          )
          aml_ctx.upload_file(name="full.log", path_or_stream=log_path)
          aml_ctx.upload_file(
              name="optimizer_state.pt",
              path_or_stream=str(
                  model_path.with_suffix(".optimizerstate")
              ),
          )
  
file_path: ptgnn/ptgnn/implementations/typilus/train.py
imports: import logging
  import random
  import torch
  from docopt import docopt
  from dpu_utils.utils import RichPath, run_and_debug
  from pathlib import Path
  from ptgnn.baseneuralmodel import ModelTrainer
  from ptgnn.baseneuralmodel.utils.amlutils import configure_logging, log_run
  from ptgnn.baseneuralmodel.utils.data import LazyDataIterable
  from ptgnn.implementations.typilus.graph2class import Graph2Class
  from ptgnn.neuralmodels.embeddings.strelementrepresentationmodel import StrElementRepresentationModel
  from ptgnn.neuralmodels.gnn import GraphNeuralNetworkModel
  from ptgnn.neuralmodels.gnn.messagepassing import GatedMessagePassingLayer, MlpMessagePassingLayer
  from ptgnn.neuralmodels.gnn.messagepassing.residuallayers import ConcatResidualLayer
label_desc: Use image node_id 3 for calling a global function with example usage: run(arguments) and returns: opt, torch