import six


def _reset_cached_subgraph(func):
    """Resets cached subgraph after execution, in case it was affected."""

    @six.wraps(func)
    def wrapper(self, *args, **kwargs):
        result = func(self, *args, **kwargs)
        self._subgraph = None
        return result

    return wrapper