black

basf · Nov 20, 2024 · 967ce17 · 967ce17
1 parent 8cfc27c
commit 967ce17
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/test_extras/test_chemprop/test_chemprop_pipeline.py b/test_extras/test_chemprop/test_chemprop_pipeline.py
@@ -308,7 +308,9 @@ def test_prediction(self) -> None:
         self.assertTrue(np.allclose(pred, pred_copy))
 
         # Test single prediction, this was causing an error before
-        single_mol_pred = regression_model.predict([molecule_net_logd_df["smiles"].iloc[0]])
+        single_mol_pred = regression_model.predict(
+            [molecule_net_logd_df["smiles"].iloc[0]]
+        )
         self.assertEqual(single_mol_pred.shape, (1,))
 
 
@@ -346,9 +348,13 @@ def test_prediction(self) -> None:
         self.assertTrue(np.allclose(proba[~nan_indices], proba_copy[~nan_indices]))
 
         # Test single prediction, this was causing an error before
-        single_mol_pred = classification_model.predict([molecule_net_bbbp_df["smiles"].iloc[0]])
+        single_mol_pred = classification_model.predict(
+            [molecule_net_bbbp_df["smiles"].iloc[0]]
+        )
         self.assertEqual(single_mol_pred.shape, (1,))
-        single_mol_proba = classification_model.predict_proba([molecule_net_bbbp_df["smiles"].iloc[0]])
+        single_mol_proba = classification_model.predict_proba(
+            [molecule_net_bbbp_df["smiles"].iloc[0]]
+        )
         self.assertEqual(single_mol_proba.shape, (1, 2))