kd.nn.train_property

kauldron.modules.train_property() → bool

is_training property.

Usage:

class MyModule(nn.Module):
  is_training = kd.nn.train_property()  # No typing annotation here !

  @nn.compact
  def __call__(self, x):
    if self.is_training:
      x = nn.Dropout(.5)(x)

    return x

The is_training property value has to be set at the top level .init() or .apply() call and will be propagated to all children.

Alternatively, the value can be set and changed inside any module with the kd.nn.set_train_property(False) context manager.

model = MyModule()
model.init(..., is_training_property=True)

Returns:

The is_training property