SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix pytorch image classification example (#14883)

Browse Source 
* Update example

* Remove skip in tests
This commit is contained in:
Mario Šaško
2021-12-22 14:42:19 +01:00
committed by GitHub
parent 7df4b90c76
commit 1045a36c1f
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
examples/pytorch/image-classification/run_image_classification.py
View File

@@ -279,12 +279,14 @@ def main():
def train_transforms(example_batch):
"""Apply _train_transforms across a batch."""
example_batch["pixel_values"] = [_train_transforms(pil_loader(f)) for f in example_batch["image_file_path"]]
example_batch["pixel_values"] = [
_train_transforms(pil_img.convert("RGB")) for pil_img in example_batch["image"]
]
return example_batch
def val_transforms(example_batch):
"""Apply _val_transforms across a batch."""
example_batch["pixel_values"] = [_val_transforms(pil_loader(f)) for f in example_batch["image_file_path"]]
example_batch["pixel_values"] = [_val_transforms(pil_img.convert("RGB")) for pil_img in example_batch["image"]]
return example_batch
if training_args.do_train: