From 14ed2f536451fcbb061c36564f42788c4acf8e3b Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Mon, 25 Jul 2022 09:36:23 -0600 Subject: [PATCH] Update model.py --- alfabet/model.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/alfabet/model.py b/alfabet/model.py index 788b4f8..22fd190 100644 --- a/alfabet/model.py +++ b/alfabet/model.py @@ -64,7 +64,9 @@ def predict(smiles_list, drop_duplicates=True, batch_size=1): tf.data.experimental.AUTOTUNE ) - bdes, bdfes = model.predict(batched_dataset) + results = model.predict(batched_dataset) + bdes = results[:, 0] + bdfes = results[:, 1] bde_df = ( pd.DataFrame(bdes.squeeze(axis=-1), index=smiles_list)