# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/18b_callback.preds.ipynb (unless otherwise specified).


from __future__ import annotations


__all__ = ['MCDropoutCallback']

# Cell
#nbdev_comment from __future__ import annotations
from ..basics import *

# Cell
class MCDropoutCallback(Callback):
    def before_validate(self):
        for m in [m for m in flatten_model(self.model) if 'dropout' in m.__class__.__name__.lower()]:
            m.train()

    def after_validate(self):
        for m in [m for m in flatten_model(self.model) if 'dropout' in m.__class__.__name__.lower()]:
            m.eval()