Update pipeline.py

This commit is contained in:
Yih-Dar SHIEH 2021-10-24 17:30:16 +00:00 committed by huggingface-web
parent c1c837b30e
commit 93bcb8d51d
1 changed files with 1 additions and 1 deletions

View File

@ -9,7 +9,7 @@ class PreTrainedPipeline():
def __init__(self, path=""): def __init__(self, path=""):
model_dir = os.path.join(path, "ckpt_epoch_3_step_6900") model_dir = path
self.model = FlaxVisionEncoderDecoderModel.from_pretrained(model_dir) self.model = FlaxVisionEncoderDecoderModel.from_pretrained(model_dir)
self.feature_extractor = ViTFeatureExtractor.from_pretrained(model_dir) self.feature_extractor = ViTFeatureExtractor.from_pretrained(model_dir)