diff --git a/package/fsrs4anki_optimizer/fsrs4anki_optimizer.py b/package/fsrs4anki_optimizer/fsrs4anki_optimizer.py
index 1a25a8a..bf4dfa8 100644
--- a/package/fsrs4anki_optimizer/fsrs4anki_optimizer.py
+++ b/package/fsrs4anki_optimizer/fsrs4anki_optimizer.py
@@ -142,7 +142,7 @@ def __init__(self, data_source: RevlogDataset, batch_size: int):
         self.generator.manual_seed(seed)
 
     def __iter__(self):
-        yield from [self.batch_indices[idx] for idx in torch.randperm(self.batch_nums, generator=self.generator).tolist()]
+        yield from (self.batch_indices[idx] for idx in torch.randperm(self.batch_nums, generator=self.generator).tolist())
 
     def __len__(self):
         return len(self.data_source)