From 9a6331b49520fd284b6929d17b4e0893805bf6d9 Mon Sep 17 00:00:00 2001 From: Ethan Harris Date: Tue, 15 Mar 2022 17:14:56 +0000 Subject: [PATCH] Fix text CI --- flash/text/question_answering/input.py | 2 +- flash/text/seq2seq/core/input.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/flash/text/question_answering/input.py b/flash/text/question_answering/input.py index 541c451e31..b1d18c371d 100644 --- a/flash/text/question_answering/input.py +++ b/flash/text/question_answering/input.py @@ -77,7 +77,7 @@ def load_data( if flash._IS_TESTING: # NOTE: must subset in this way to return a Dataset - hf_dataset = [sample for sample in hf_dataset.select(range(40), keep_in_memory=True)] + hf_dataset = [sample for sample in hf_dataset][:40] return hf_dataset diff --git a/flash/text/seq2seq/core/input.py b/flash/text/seq2seq/core/input.py index 0f935d49aa..88d14adde6 100644 --- a/flash/text/seq2seq/core/input.py +++ b/flash/text/seq2seq/core/input.py @@ -44,7 +44,7 @@ def load_data( if flash._IS_TESTING: # NOTE: must subset in this way to return a Dataset - hf_dataset = [sample for sample in hf_dataset.select(range(40), keep_in_memory=True)] + hf_dataset = [sample for sample in hf_dataset][:40] return hf_dataset