def METHOD_NAME(self, X):
    raw = self._wrapped_model.METHOD_NAME(X)
    if isinstance(raw, (np.ndarray, pd.DataFrame)):
        s_X = lale.datasets.data_schemas.to_schema(X)
        s_result = self.transform_schema(s_X)
        result = lale.datasets.data_schemas.add_schema(raw, s_result, recalc=True)
    else:
        result = raw
    return result