lunny/keras
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix demo

Browse Source
This commit is contained in:
Francois Chollet 2023-05-03 11:11:30 -07:00
parent 5a9fc53a95
commit 1e3e94cf3c
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
demo_custom_jax_workflow.py
View File

@ -113,6 +113,7 @@ for data in dataset:
print("Loss:", loss)
# Post-processing model state update
trainable_variables, non_trainable_variables, optimizer_variables = state
for variable, value in zip(model.trainable_variables, trainable_variables):
variable.assign(value)
for variable, value in zip(