search

deepmodeling / Uni-Mol

Official Repository for the Uni-Mol Series Methods
MIT License
726 stars 127 forks source link

unimol_tools do not make use of gpu when infer #212

Open PoloWitty opened 7 months ago

PoloWitty commented 7 months ago

The code implemented in unimol_tools did not pass the cuda variable to Trainer, so the gpu won't be used.

The related code is here.

Should be changed like this:

self.trainer = Trainer(task='repr', cuda=self.device)
emotionor commented 7 months ago

Thank you for bringing this issue to our attention and for providing a solution.

Your contribution is greatly appreciated.

We will review the fix and incorporate it if appropriate. Thanks again for your help!