fix batch tokenisation

Author: sfluegel05

chebai/preprocessing/datasets/pubchem.py

@@ -285,18 +285,19 @@ def _tokenize_batched(self, data):
                     b
                     for b in batch
                     if b["features"] is not None
-                    and (self.n_token_limit is None
-                    or len(b["features"]) <= self.n_token_limit)
                 ]
+                if self.n_token_limit is not None:
+                    batch = [b for b in batch if len(b["features"]) <= self.n_token_limit]
                 yield batch
                 batch = []
         print("Saving final batch")
         batch = [
             b
             for b in batch
             if b["features"] is not None
-            and (self.n_token_limit is None or len(b["features"]) <= self.n_token_limit)
         ]
+        if self.n_token_limit is not None:
+            batch = [b for b in batch if len(b["features"]) <= self.n_token_limit]
         yield batch
 
     def setup_processed(self):