code: def eval_step(network, data):
      inputs = {
          "obs": data[:, :-1],
          "target": data[:, 1:],
      }
  
      out = network.eval(inputs)
      loss = out["loss"]
  
      return np.array(loss).mean()
  
file_path: mesh-transformer-jax/device_train.py
imports: import argparse
  import json
  import time
  import jax
  import numpy
  import optax
  import wandb
  from tqdm import tqdm
  from mesh_transformer import util
  from mesh_transformer.checkpoint import read_ckpt, write_ckpt
  from mesh_transformer.transformer_shard import CausalTransformer
  from tfrecord_loader import TFRecordNewInputs
  from smart_open import open
  from google.cloud import storage
  from google.cloud.exceptions import NotFound
  from mesh_transformer.util import clip_by_global_norm, additive_weight_decay
label_desc: Use image node_id 4 for calling a global function with example usage: eval_step(network, data) and returns: np