def with_graph(f):
    """Call a function with self.graph.as_default() context."""
    def wrapper(*args, **kwargs):
        # the first argument is always the model instance
        model = args[0]
        with model.graph.as_default():
            return f(*args, **kwargs)
    return wrapper