TrainMode

Puts the model or node into training mode (enabling dropout, batch norm updating).

Abstract Signature:

TrainMode(node)

PyTorchJAX (Core)KerasTensorFlowApple MLXFlax NNX

PyTorch

API: —

Strategy: Macro '{node}.train()'

JAX (Core)

API: —

Strategy: Custom / Partial

Keras

API: —

Strategy: Macro 'setattr({node}, 'trainable', True) or {node}'

TensorFlow

API: —

Strategy: Custom / Partial

Apple MLX

API: —

Strategy: Custom / Partial

Flax NNX

API: flax.nnx.train_mode

Strategy: Direct Mapping

Official Docs ↗