code: def load_and_predict(path: str, idx: int = 0):
      """
      Load the saved msgpack model and make predictions.
      We will run prediction on test MNIST dataset
      """
      if not os.path.exists(path):
          raise FileNotFoundError(f"File not found at {path}")
      with open(path, "rb") as f:
          state_dict = serialization.from_bytes(CNN, f.read())
      cnn = CNN()
      _, test_ds = get_datasets()
      # ensure that all arrays are restored as jnp.ndarray
      # NOTE: This is to prevent a bug this will be fixed in Flax >= v0.3.4:
      # https://github.com/google/flax/issues/1261
      if pkg_version_info("flax") < (0, 3, 4):
          state_dict = jax.tree_util.tree_map(jnp.ndarray, state_dict)
      # jit it !
      logits = jax.jit(
          lambda x: cnn.apply({"params": state_dict["params"]}, x)
      )(test_ds["image"])
      return logits[idx].argmax()
  
imports: from __future__ import annotations
  import argparse
  import os
  import typing
  from typing import TYPE_CHECKING
  import attrs
  import cattrs
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
  from flax import linen
  from flax import serialization
  from flax.metrics import tensorboard
  from flax.training import train_state
  import bentoml
  from bentoml._internal.utils.pkg import pkg_version_info
label_desc: Use image node_id 7 for calling a global function with example usage: load_and_predict(path, idx) and returns: logits