Logits dont match dim(design_residues) when using --pocket_def_ligand

[Lukasra00]

Hi Hannes,

Thanks for this amazing tool, you built!

Problem: I am using a .csv to run the inferences.py script. I want to jointly generate logits for the three design_residues (as specified in the csv) and dock the ligand into the enzyme. However, the logits are only of dimensionality (num_inference, amino_acids=21). The dimensionality of the logits should in my understanding be (num_inference, #design_residues, amino_acids=21).

In contrast: When I define the pocket not using --pocket_def_ligand, but --pocket_def_residues the logit dimensionality is fine and matches the expected dim (num_inference, #design_residues, amino_acids=21).

Recreating the problem:
.csv format: ligand,smiles,protein,design_residues,pocket_def_ligand,pocket_def_residues,pocket_def_center ,C1=CC2=C(C=CN2)C=C1I O=Cc1c(O)c(C)ncc1COP(O)(O)=O,../enzymeoracle/data/structure/PfTrpB-5iodo.pdb,"A165, A183, A301",../enzymeoracle/data/structure/5iodo_PLP.sdf,,

CLI: CUDA_VISIBLE_DEVICES="" python -m inference --num_inference 1 --out_dir ../enzymeoracle/outputs/flowsite/output --csv_file dummy_inference.csv --batch_size 16 --checkpoint pocket_gen/b1ribx1a/checkpoints/best.ckpt --run_test --run_name inference1 --layer_norm --fake_constant_dur 100000 --fake_decay_dur 10 --fake_ratio_start 0.2 --fake_ratio_end 0.2 --residue_loss_weight 0.2 --use_tfn --use_inv --time_condition_tfn --correct_time_condition --time_condition_inv --time_condition_repeat --flow_matching --flow_matching_sigma 0.5 --prior_scale 1 --num_angle_pred 5 --ns 32 --nv 8 --batch_norm --self_condition_x --self_condition_inv --no_tfn_self_condition_inv --self_fancy_init --pocket_residue_cutoff 12

Any help / clarification on how I can obtain logits for designed_residues is appreciated.
Thanks!

Best,
Lukas