Update src/datasets.py

Disable second filtering by number of atoms

src/datasets.py

@@ -218,10 +218,10 @@ def collate(batch):
     out = {}
 
     # Filter out big molecules
-    if 'pocket_mask' not in batch[0].keys():
-        batch = [data for data in batch if data['num_atoms'] <= 50]
-    else:
-        batch = [data for data in batch if data['num_atoms'] <= 1000]
+    # if 'pocket_mask' not in batch[0].keys():
+    #    batch = [data for data in batch if data['num_atoms'] <= 50]
+    # else:
+    # batch = [data for data in batch if data['num_atoms'] <= 1000]
 
     for i, data in enumerate(batch):
         for key, value in data.items():