Revert "Update model.py"

This reverts commit 14ed2f5.
NREL · Jul 25, 2022 · 8ecbe4b · 8ecbe4b
1 parent 14ed2f5
commit 8ecbe4b
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/alfabet/model.py b/alfabet/model.py
@@ -64,9 +64,7 @@ def predict(smiles_list, drop_duplicates=True, batch_size=1):
         tf.data.experimental.AUTOTUNE
     )
 
-    results = model.predict(batched_dataset)
-    bdes = results[:, 0]
-    bdfes = results[:, 1]
+    bdes, bdfes = model.predict(batched_dataset)
 
     bde_df = (
         pd.DataFrame(bdes.squeeze(axis=-1), index=smiles_list)