flax-community
/

vit-gpt2

Model card Files Files and versions Metrics Training metrics Community

ydshieh commited on Jul 17, 2021

Commit

790ecb1

•

1 Parent(s): 7a16b57

update generate.py

Files changed (1) hide show

generate.py +1 -1

generate.py CHANGED Viewed

@@ -68,7 +68,7 @@ if __name__ == '__main__':
     print(f'caption: {caption}')
     ds = datasets.load_dataset('./coco_dataset_script.py', data_dir='/home/33611/caption/')
-    ds = ds['train']
     ds = ds.select(range(100))
     predictions = []

     print(f'caption: {caption}')
     ds = datasets.load_dataset('./coco_dataset_script.py', data_dir='/home/33611/caption/')
+    ds = ds['validation']
     ds = ds.select(range(100))
     predictions = []