code: def train_func(config):
      """Your training function that will be launched on each worker."""
  
      # Unpack training configs
      set_seed(config["seed"])
      num_epochs = config["num_epochs"]
      train_batch_size = config["train_batch_size"]
      eval_batch_size = config["eval_batch_size"]
  
      # Instantiate the Model
      model = AutoModelForSequenceClassification.from_pretrained(
          "bert-base-cased", return_dict=True
      )
  
      # Prepare Ray Data Loaders
      # ====================================================
      train_ds = ray.train.get_dataset_shard("train")
      eval_ds = ray.train.get_dataset_shard("validation")
  
      tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
  
      def collate_fn(batch):
          outputs = tokenizer(
              list(batch["sentence1"]),
              list(batch["sentence2"]),
              truncation=True,
              padding="longest",
              return_tensors="pt",
          )
          outputs["labels"] = torch.LongTensor(batch["label"])
          return outputs
  
      train_dataloader = train_ds.iter_torch_batches(
          batch_size=train_batch_size, collate_fn=collate_fn
      )
      eval_dataloader = eval_ds.iter_torch_batches(
          batch_size=eval_batch_size, collate_fn=collate_fn
      )
      # ====================================================
  
      # Initialize DeepSpeed Engine
      model, optimizer, _, lr_scheduler = deepspeed.initialize(
          model=model,
          model_parameters=model.parameters(),
          config=deepspeed_config,
      )
      device = get_accelerator().device_name(model.local_rank)
  
      # Initialize Evaluation Metrics
      f1 = BinaryF1Score().to(device)
      accuracy = BinaryAccuracy().to(device)
  
      for epoch in range(num_epochs):
          # Training
          model.train()
          for batch in train_dataloader:
              batch = {k: v.to(device) for k, v in batch.items()}
              outputs = model(**batch)
              loss = outputs.loss
              model.backward(loss)
              optimizer.step()
              lr_scheduler.step()
              optimizer.zero_grad()
  
          # Evaluation
          model.eval()
          for batch in eval_dataloader:
              batch = {k: v.to(device) for k, v in batch.items()}
              with torch.no_grad():
                  outputs = model(**batch)
              predictions = outputs.logits.argmax(dim=-1)
  
              f1.update(predictions, batch["labels"])
              accuracy.update(predictions, batch["labels"])
  
          # torchmetrics will aggregate the metrics across all workers
          eval_metric = {
              "f1": f1.compute().item(),
              "accuracy": accuracy.compute().item(),
          }
          f1.reset()
          accuracy.reset()
  
          if model.global_rank == 0:
              print(f"epoch {epoch}:", eval_metric)
  
          # Report checkpoint and metrics to Ray Train
          # ==============================================================
          with TemporaryDirectory() as tmpdir:
              # Each worker saves its own checkpoint shard
              model.save_checkpoint(tmpdir)
  
              # Ensure all workers finished saving their checkpoint shard
              torch.distributed.barrier()
  
              # Report checkpoint shards from each worker in parallel
              ray.train.report(
                  metrics=eval_metric,
                  checkpoint=Checkpoint.from_directory(tmpdir),
              )
  
file_path: ray/python/ray/train/examples/deepspeed/deepspeed_torch_trainer.py
imports: from tempfile import TemporaryDirectory
  import deepspeed
  import torch
  from datasets import load_dataset
  from deepspeed.accelerator import get_accelerator
  from torchmetrics.classification import BinaryAccuracy, BinaryF1Score
  from transformers import AutoModelForSequenceClassification, AutoTokenizer, set_seed
  import ray
  import ray.train
  from ray.train import Checkpoint, DataConfig, ScalingConfig
  from ray.train.torch import TorchTrainer
label_desc: Use image node_id 1 for calling a global function with example usage: train_func(config) and returns: outputs