Update model.py

NREL · Jul 25, 2022 · 14ed2f5 · 14ed2f5
1 parent 21c2419
commit 14ed2f5
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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)