Merge pull request #691 from chemprop/interpret-num-workers

Pass num_workers to MoleculeDataLoader during interpretation
chemprop · Mar 1, 2024 · 0f0da06 · 0f0da06
2 parents 1785627 + ee99bd7
commit 0f0da06
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/chemprop/interpret.py b/chemprop/interpret.py
@@ -56,7 +56,7 @@ def __call__(self, smiles: List[str], batch_size: int = 500) -> List[List[float]
         if self.train_args.bond_descriptor_scaling and self.args.bond_descriptors_size > 0:
             test_data.normalize_features(self.bond_descriptor_scaler, scale_bond_descriptors=True)
 
-        test_data_loader = MoleculeDataLoader(dataset=test_data, batch_size=batch_size)
+        test_data_loader = MoleculeDataLoader(dataset=test_data, batch_size=batch_size, num_workers=self.args.num_workers)
 
         sum_preds = []
         for model in self.checkpoints: