search

arneschneuing / DiffSBDD

A Euclidean diffusion model for structure-based drug design.
MIT License
324 stars 72 forks source link

cuda memory usage #7

Closed BL-Lac149597870 closed 1 year ago

BL-Lac149597870 commented 1 year ago

Hi @arneschneuing, when training the model with batchsize=1, I get the following error:

image

It seems there is somethig wrong with the cuda memory (caused by pytorch-lightening?). I don't think the training process would need about 40GB cuda memory with batchsize=1...

arneschneuing commented 1 year ago

Hi @BL-Lac149597870, this can indeed happen if you train a full-atom model with a fully-connected graph because a huge number of edges will be created. Could you please try again with a distance threshold for edges? It can be specified in the config file (in the preprint we use edge_cutoff: 7.0). I hope this will solve the issue. Best, Arne

BL-Lac149597870 commented 1 year ago

Thanks a lot @arneschneuing , that works for me. Thank you very much and really a nice work (I saw the comments from ICLR2023 and I think ICLR reviewers are too harsh) !