fix fill_diag (#219)

deepmodeling · May 6, 2024 · 31df659 · 31df659
1 parent 7394647
commit 31df659
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/unimol/unimol/utils/docking_utils.py b/unimol/unimol/utils/docking_utils.py
@@ -181,7 +181,7 @@ def docking_data_pre(raw_data_path, predict_path):
             distance_predict = distance_predict.numpy().astype(np.float32)
             holo_distance_predict = holo_distance_predict.numpy().astype(np.float32)
             # Fill diagonal with 0, issue with the model not learning to predict 0 distance
-            holo_distance_predict = np.fill_diagonal(holo_distance_predict, 0)
+            np.fill_diagonal(holo_distance_predict, 0)
             #
             holo_coords = holo_coordinates.numpy().astype(np.float32)