def check_norm_state(modules, train_state):
    """Check if norm layer is in correct train state."""
    for mod in modules:
        if isinstance(mod, _BatchNorm):
            if mod.training != train_state:
                return False
    return True
Check if norm layer is in correct train state.
tests/test_models/test_backbones/test_tnt.py