Issue running full atom joint model

[amitkadan]

Hi there, I am working on reproducing the results reported in the paper on the crossdocked dataset, and was able to successfully do so when running test.py with the crossdocked_fullatom_cond and the crossdocked_ca_joint model, i.e., by running
python test.py crossdocked_ca_joint.ckpt --test_dir <crossdocked_dir>/processed_noH/test/ --outdir <output_dir> --sanitize --timesteps 50 --n_nodes_bias 5
However, when I try to run the crossdocked_fullatom_joint model I get,

0%|                                                                                                        | 0/100
[00:00<?, ?it/s]/shared/akadan/DiffSBDD/test.py:172: UserWarning: Attempt 1/10 failed with error: 'NaN detected in EGNN output'. Trying again...
...
RuntimeError: Maximum number of retries exceeded

I tried running generate_ligands.py with this same model and parameters, and noticed that this error always arises after 5 diffusion steps. Am I doing something wrong when running this model? It seems be the right parameters as specified in the paper (i.e. 50 diffusion steps with 10 resampling steps).

Thanks in advance,
Amit