python - モデルのトレーニング後に Trax で予測モードに切り替える方法は?

Question

ディープ ラーニング モデルを構築するための Trax ライブラリに慣れてきましたが、答えが見つからない 1 つの質問は、モデルのトレーニングが完了した後に "train" モードから "eval" モードに切り替える方法です。

トランスフォーマーを使用したこの例を検討してください ( https://github.com/google/trax#1-run-a-pre-trained-transformerから変更)

import trax

model = trax.models.Transformer(
    input_vocab_size=33300,
    d_model=512, d_ff=2048,
    n_heads=8, n_encoder_layers=6, n_decoder_layers=6,
    max_len=2048, mode='train')

# do the training

トレーニング後、モードを「予測」に切り替えるにはどうすればよいですか?

(同じ例に従って) 頭に浮かぶアイデアの 1 つは、モデルをトレーニングしてから保存し、モデルを再初期化することですが、今回は を使用してmode='train'、実行して重みをロードしますmodel.init_from_file(file)。

モデルを再初期化せずに直接行う方法はありますか?