def _get_model_or_module(crypten_model):
    """
    Returns `Module` if model contains only one module. Otherwise returns model.
    """
    num_modules = len(list(crypten_model.modules()))
    if num_modules == 1:
        for crypten_module in crypten_model.modules():
            return crypten_module
    return crypten_model