code: def load_model() -> "MockModel":
      """Load MNIST model.
  
      Returns
      -------
      MnistModel
      """
      path = MODEL_PATH
      saved_path = MODEL_SAVED_PATH  # Different path because tf saves the model in three files, none named as MODEL_PATH
  
      def create_model():
          """Create a new model."""
          return keras.models.Sequential(
              [
                  keras.layers.Flatten(input_shape=(28, 28, 1)),
                  keras.layers.Dense(128, activation="relu"),
                  keras.layers.Dropout(0.2),
                  keras.layers.Dense(10),
              ]
          )
  
      def add_softmax(model: keras.models.Sequential):
          """Add softmax layer to model."""
          return keras.Sequential([model, keras.layers.Softmax()])
  
      if saved_path.exists():
          model = create_model()
          model.load_weights(path).expect_partial()
          model = add_softmax(model)
          return MockModel(model)
  
      model = create_model()
      loss_fn = keras.losses.SparseCategoricalCrossentropy(
          from_logits=True
      )
      model.compile(
          optimizer="adam", loss=loss_fn, metrics=["accuracy"]
      )
  
      x_train, y_train = load_mnist_data(train=True)
  
      model.fit(x_train, y_train, epochs=2)
      del x_train, y_train
  
      model.save_weights(path)
  
      # Change output to softmax for probabilities
      model = add_softmax(model)
  
      if not path.parent.exists():
          path.parent.mkdir()
  
      return MockModel(model)
  
file_path: deepchecks/deepchecks/vision/datasets/classification/mnist_tensorflow.py
imports: import logging
  import pathlib
  import pickle
  import typing
  import numpy
  from deepchecks.vision.utils.test_utils import hash_image
  from deepchecks.vision.vision_data import VisionData
label_desc: Use image node_id 4 for calling a global function with example usage: load_model() and returns: MockModel, keras, keras, MockModel