code: def evaluate_model(
      model: nn.Module,
      test_dataloader: Union[DataLoader, List[DataLoader]],
      num_labels: int,
      task_name: str,
      eval_splits: List[str],
      coordinator: DistCoordinator,
  ):
      metric = evaluate.load(
          "glue",
          task_name,
          process_id=coordinator.rank,
          num_process=coordinator.world_size,
      )
      model.eval()
  
      def evaluate_subset(dataloader: DataLoader):
          accum_loss = torch.zeros(
              1, device=torch.cuda.current_device()
          )
          for batch in dataloader:
              batch = to_device(batch, torch.cuda.current_device())
              outputs = model(**batch)
              val_loss, logits = outputs[:2]
              accum_loss.add_(val_loss)
  
              if num_labels > 1:
                  preds = torch.argmax(logits, axis=1)
              elif num_labels == 1:
                  preds = logits.squeeze()
  
              labels = batch["labels"]
              metric.add_batch(predictions=preds, references=labels)
  
          results = metric.compute()
          if coordinator.is_master():
              results["loss"] = accum_loss.item() / (
                  len(dataloader) * dataloader.batch_size
              )
          return results
  
      if isinstance(test_dataloader, DataLoader):
          return evaluate_subset(test_dataloader)
      else:
          assert len(test_dataloader) == len(eval_splits)
          final_results = {}
          for split, sub_loader in zip(eval_splits, test_dataloader):
              results = evaluate_subset(sub_loader)
              final_results.update(
                  {f"{k}_{split}": v for k, v in results.items()}
              )
          return final_results
  
file_path: colossalai/colossalai/shardformer/examples/convergence_benchmark.py
imports: import argparse
  import math
  from typing import Any, List, Union
  import evaluate
  import torch
  import torch.distributed
  from data import GLUEDataBuilder
  from torch import nn
  from torch.optim import Adam, Optimizer
  from torch.utils._pytree import tree_map
  from torch.utils.data import DataLoader
  from tqdm import tqdm
  from transformers import BertConfig, BertForSequenceClassification, get_linear_schedule_with_warmup
  import colossalai
  from colossalai.cluster import DistCoordinator
  from colossalai.shardformer import ShardConfig, ShardFormer
label_desc: Use image node_id 4 for calling a global function with example usage: evaluate_model(model, test_dataloader, num_labels, task_name, eval_splits, coordinator) and returns: results, evaluate_subset, final_results