In Compose, replace lambda function with named function (#151)

This prevents the following error on Windows (when using
a multi-process DataLoader, for example):

AttributeError: Can't pickle local object '_transform.<locals>.<lambda>'
This commit is contained in:
Gianluca Gippetto 2021-09-24 03:42:20 +02:00 committed by GitHub
parent 3b473b0e68
commit c13005fd42
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 1 deletions

View File

@ -68,11 +68,15 @@ def _download(url: str, root: str):
return download_target
def _convert_image_to_rgb(image):
return image.convert("RGB")
def _transform(n_px):
return Compose([
Resize(n_px, interpolation=BICUBIC),
CenterCrop(n_px),
lambda image: image.convert("RGB"),
_convert_image_to_rgb,
ToTensor(),
Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
])