bugfix

mokawi · Mar 16, 2021 · c241b97 · c241b97
1 parent 9289bcf
commit c241b97
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/clctm.py b/clctm.py
@@ -131,7 +131,7 @@ def _get_indices(self, doc_idx):
         if isinstance(doc_idx, int):
             return list(range(*self.doc_rng(doc_idx)))
         else:
-            return list(chain(*(range(*self.doc_rng[d] for d in doc_idx))))
+            return list(chain(*(range(*self.doc_rng[d]) for d in doc_idx)))
 
     def _get_mask(self, doc_idx):
         r = np.full(len(self.input_ids), False)