generatebio / chroma

A generative model for programmable protein design
Apache License 2.0
627 stars 74 forks source link

Expected all tensors to be on the same device #10

Closed QUEST2179 closed 7 months ago

QUEST2179 commented 7 months ago

Dear developer,

please give a hand on this issue, Thanks!

this example look fine from chroma import Chroma chroma = Chroma() protein = chroma.sample(chain_lengths=[200]) protein.to("sample.cif")

But try running the example from chroma import Protein, Chroma chroma = Chroma() protein = Protein('1GFP') protein = chroma.design(protein) protein.to("1GFP-redesign.cif")

gave error messages: Traceback (most recent call last): File "testDesignwithChroma.py", line 21, in protein = chroma.design(protein) File "C:\Users\project\structure\chroma\chroma\models\chroma.py", line 528, in design X_sample, Ssample, = self.design_network.sample( File "C:\Users\miniconda3\envs\py8\lib\site-packages\torch\autograd\grad_mode.py", line 27, in decorate_context return func(*args, kwargs) File "C:\Users\project\structure\chroma\chroma\data\xcs.py", line 114, in new_func return func(*args, *kwargs) File "C:\Users\project\structure\chroma\chroma\models\graph_design.py", line 826, in sample node_h, edge_h, edge_idx, mask_i, mask_ij = self.encode(X, C, t=t) File "C:\Users\project\structure\chroma\chroma\data\xcs.py", line 114, in new_func return func(args, kwargs) File "C:\Users\project\structure\chroma\chroma\models\graph_design.py", line 536, in encode node_h, edge_h, edge_idx, mask_i, mask_ij = self.encoder( File "C:\Users\miniconda3\envs\py8\lib\site-packages\torch\nn\modules\module.py", line 1130, in _call_impl return forward_call(*input, *kwargs) File "C:\Users\project\structure\chroma\chroma\data\xcs.py", line 114, in new_func return func(args, kwargs) File "C:\Users\project\structure\chroma\chroma\models\graph_design.py", line 1237, in forward node_h, edge_h, edge_idx, mask_i, mask_ij = self._checkpoint( File "C:\Users\project\structure\chroma\chroma\models\graph_design.py", line 1253, in _checkpoint return module(args) File "C:\Users\miniconda3\envs\py8\lib\site-packages\torch\nn\modules\module.py", line 1130, in _call_impl return forward_call(input, kwargs) File "C:\Users\project\structure\chroma\chroma\layers\structure\protein_graph.py", line 223, in forward node_h_l = node_h_l - self.getattr(f"nodemeans{i}") RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

wujiewang commented 7 months ago

Hey, Protein object needs to be sent to "cuda" too

you just need to do

protein = Protein('1GFP', device='cuda')