search

DeepChainBio / bio-transformers

bio-transformers is a wrapper on top of the ESM/Protbert model, trained on millions on proteins and used to predict embeddings.
https://bio-transformers.readthedocs.io/en/latest/getting_started/install.html
Apache License 2.0
143 stars 31 forks source link

add multi-gpu support #2

Closed delfosseaurelien closed 3 years ago

delfosseaurelien commented 3 years ago

allow the use of torch multi-gpu for inference use of nn.DataParallel or ray to make the inference on multiple batches

mathpluscode commented 3 years ago

I believe nn.DataParallele is the recommended solution considering the native support of mixed precision https://pytorch.org/docs/stable/notes/amp_examples.html#working-with-multiple-gpus