Как определить переменную в TensorFlow

Я застрял в определении переменных.

Мой код:

import tensorflow as tf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten, Conv2D, MaxPooling2D

mnist_data = tf.keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist_data.load_data()``


def scale_mnist_data(train_images, test_images):

return (train_images / 255, test_images / 255)


def train_model(model, scaled_train_images, train_labels):
scaled_train_images, scaled_test_images = scale_mnist_data(train_images, test_images)

До этого момента код работает гладко, но здесь

scaled_train_images = scaled_train_images[..., np.newaxis]
scaled_test_images = scaled_test_images[..., np.newaxis]

Я получаю эту ошибку:

NameError                                 Traceback (most recent call last)
<ipython-input-5-7e4c845d2449> in <module>
  1 # Add a dummy channel dimension
  2 
----> 3 scaled_train_images = scaled_train_images[..., np.newaxis]
  4 scaled_test_images = scaled_test_images[..., np.newaxis]

NameError: name 'scaled_train_images' is not defined

Интересно, можно ли вставить этот код def train_model(model, scaled_train_images, train_labels):. Но здесь я снова столкнулся с похожими проблемами, такими как история, кадр и некоторые другие переменные, которые невозможно определить.

К вашему сведению: я пытаюсь запустить свой код на курсе Coursera для Имперского лондонского колледжа: Начало работы с TensorFlow 2.

Я новичок, не знаю Python.