Code Error, page 246, Knowledge Distillation

[dia2teen]

---

KeyError Traceback (most recent call last)
in <cell line: 5>()
26 vars = s_model_2.trainable_variables
27 grad = tape.gradient(student_loss, vars)
---> 28 opt.apply_gradients(zip(grad, vars))
29
30 print("에포크 {}".format(e))

12 frames
/usr/local/lib/python3.10/dist-packages/keras/src/optimizers/optimizer.py in _update_step(self, gradient, variable)
231 return
232 if self._var_key(variable) not in self._index_dict:
--> 233 raise KeyError(
234 f"The optimizer cannot recognize variable {variable.name}. "
235 "This usually means you are trying to call the optimizer to "

KeyError: in user code:

File "/usr/local/lib/python3.10/dist-packages/keras/src/optimizers/optimizer.py", line 225, in _update_step_xla  *
    return self._update_step(gradient, variable)
File "/usr/local/lib/python3.10/dist-packages/keras/src/optimizers/optimizer.py", line 233, in _update_step  **
    raise KeyError(

KeyError: 'The optimizer cannot recognize variable dense_1/kernel:0. This usually means you are trying to call the optimizer to update different parts of the model separately. Please call `optimizer.build(variables)` with the full list of trainable variables before the training loop or use legacy optimizer `tf.keras.optimizers.legacy.Adam.'

page 246쪽에 있는 코드를 실행하면 Key Error가 발생합니다.
optimizer.build(variables)를 통해 변수의 리스트를 만드라는 건데, 어떻게 해결해야 할까요?