[flax:examples:sst2] Fix notebook error.#5122

copybara-service[bot] merged 1 commit intomaingoogle/flax:mainfrom

test_839349552google/flax:test_839349552Copy head branch name to clipboard