SetMode

Sets the execution mode (train/eval) for a module.

Abstract Signature:

SetMode(node, training: bool)

PyTorchKerasFlax NNX

PyTorch

API: train

Strategy: Plugin (torch_mode_switch)

Keras

API: set_training_mode

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: flax.nnx.set_mode

Strategy: Direct Mapping

Official Docs ↗