code: def to_device(self, batch, keys=["input_ids", "labels"]):
      """Move batch of data into device memory."""
      return {
          k: v.to(
              device=self.setup["device"],
              dtype=torch.long,
              non_blocking=True,
          )
          for k, v in batch.items()
          if k in keys  # Add more keywords here if needed
      }
  
file_path: bert-ultrafast/training/cramming/backend/deepspeed_integration.py
imports: import torch
  import os
  import json
  from functools import partial
  import logging
  from omegaconf import OmegaConf
  from .utils import group_parameters, prepare_pretraining_dataloader
  from .optimizers import get_schedule_fn
label_desc: Use image node_id 5 for calling a global function with example usage: to_device(batch, keys) and returns: unknown