Do not keep duplicate predict split in memory

artemisp · Mar 25, 2024 · 8e1dfc5 · 8e1dfc5
1 parent 1eb8553
commit 8e1dfc5
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/parallelm/data/pl_dataloaders.py b/parallelm/data/pl_dataloaders.py
@@ -263,6 +263,7 @@ def convert_label(examples):
 
         if kwargs.get('predict_split', None):
             raw_dataset['predict'] =  raw_dataset[kwargs['predict_split']]
+            del raw_dataset[kwargs['predict_split']]
         return raw_dataset
 
     def filter_long_sequences(self, **kwargs):