INNER CODE UNIT · Python
predict_coord
QizhiPei/FABind · FABind/fabind/fabind_inference.py:303
predict_coord = predict_coord - predict_coord.mean(dim=0).reshape(1, 3) + com_coord_pred_center_i
com_coord_pred[i_mask] = predict_coord
else:
predict_coord, loss, rmsd = post_optimize_compound_coords(
reference_compound_coords=com_coord_i.to(post_optim_device),
predict_compound_coords=com_coord_pred_i.to(post_optim_device),
# LAS_edge_index=(data[i]['complex', 'LAS', 'complex'].edge_index - data[i]['complex', 'LAS', 'complex'].edge_index.min()).to(post_optim_device),
LAS_edge_index=LAS_tmp[i].to(post_optim_device),
mode=args.post_optim_mode,
total_epoch=args.post_optim_epoch,
)
predict_coord = predict_coord.to(accelerator.device)
predict_coord = predict_coord - predict_coord.mean(dim=0).reshape(1, 3) + com_coord_pred_center_i
com_coord_pred[i_mask] = predict_coord
com_coord_pred_per_sample_list.append(com_coord_pred[i_mask])
com_coord_per_sample_list.append(com_coord_i)
com_coord_offset_per_sample_list.append(data[i].coord_offset)