import paddle

from doc_dewarp.GeoTr import GeoTr

model_path = "model/dewarp_model/best.ckpt"
checkpoint = paddle.load(model_path)
state_dict = checkpoint["model"]
DEWARP = GeoTr()
DEWARP.set_state_dict(state_dict)
DEWARP.eval()