probly.representer.sampler.flax.register_forced_train_mode¶ probly.representer.sampler.flax.register_forced_train_mode(cls: LazyType) → None[source]¶ Register a class to be forced into train mode during sampling. This enables Monte Carlo sampling techniques like MC Dropout [GG16a].