From 86dddc25095e67e7983ec9eb8bd40e97ff3bdf17 Mon Sep 17 00:00:00 2001 From: Parmeet Singh Bhatia Date: Tue, 11 Jan 2022 13:31:20 -0500 Subject: [PATCH] Fix filter logic --- torchtext/datasets/amazonreviewpolarity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchtext/datasets/amazonreviewpolarity.py b/torchtext/datasets/amazonreviewpolarity.py index 6a0497c3a1..e2b99e1a43 100644 --- a/torchtext/datasets/amazonreviewpolarity.py +++ b/torchtext/datasets/amazonreviewpolarity.py @@ -46,5 +46,5 @@ def AmazonReviewPolarity(root: str, split: Union[Tuple[str], str]): cache_dp = GDriveReader(cache_dp).end_caching(mode="wb", same_filepath_fn=True) cache_dp = FileOpener(cache_dp, mode="b") extracted_files = cache_dp.read_from_tar() - filter_extracted_files = extracted_files.filter(lambda x: split in x[0]) + filter_extracted_files = extracted_files.filter(lambda x: _EXTRACTED_FILES[split] in x[0]) return filter_extracted_files.parse_csv().map(fn=lambda t: (int(t[0]), ' '.join(t[1:])))