code: def worker_init(trainer: DistributedModelTrainer, rank, world_isze):
      try:
          from azureml.core.run import Run
  
          aml_ctx = Run.get_context()
      except:
          aml_ctx = None
  
      log_path = configure_logging(aml_ctx, rank=rank)
  
      trainer.register_train_epoch_end_hook(
          partial(log_run_lambda, aml_ctx, "train-" + str(rank))
      )
      trainer.register_validation_epoch_end_hook(
          partial(log_run_lambda, aml_ctx, "valid-" + str(rank))
      )
  
      def upload_hook(model, nn, epoch, metrics):
          aml_ctx.upload_file(
              name="model.pkl.gz",
              path_or_stream=str(trainer._checkpoint_location),
          )
          aml_ctx.upload_file(name="full.log", path_or_stream=log_path)
          aml_ctx.upload_file(
              name="optimizer_state.pt",
              path_or_stream=str(
                  trainer._checkpoint_location.with_suffix(
                      ".optimizerstate"
                  )
              ),
          )
  
      if rank == 0 and aml_ctx is not None:
          trainer.register_epoch_improved_end_hook(upload_hook)
  
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 4 for calling a global function with example usage: worker_init(trainer, rank, world_isze) without return types