class EstimatorAdapter: def __init__(self, estimator): self.estimator = estimator def predict(self, batch): return self.estimator(batch)