code: def load_from_folder(
      path: RichPath, shuffle: bool, rank: int, world_size
  ):
      all_files = [
          p
          for i, p in enumerate(
              path.get_filtered_files_in_dir("*.jsonl.gz")
          )
          if i % world_size == rank
      ]
  
      if shuffle:
          random.shuffle(all_files)
      for file in all_files:
          yield from file.read_as_jsonl()
  
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 1 for calling a global function with example usage: load_from_folder(path, shuffle, rank, world_size) without return types