RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

[linuxonly801]

Hello, I tried the demo code today, but the following RuntimeError occurs. Please help to fix this, THX.

/data/personal/yankai/packages/anaconda3/envs/igfold/lib/python3.7/site-packages/Bio/pairwise2.py:283: BiopythonDeprecationWarning: Bio.pairwise2 has been deprecated, and we intend to remove it in a future release of Biopython. As an alternative, please consider using Bio.Align.PairwiseAligner as a replacement, and contact the Biopython developers if you still need the Bio.pairwise2 module. BiopythonDeprecationWarning,

The code, data, and weights for this work are made available for non-commercial use 
(including at commercial entities) under the terms of the JHU Academic Software License 
Agreement. For commercial inquiries, please contact dmalon11[at]jhu.edu.
License: https://github.com/Graylab/IgFold/blob/main/LICENSE.md

Loading 4 IgFold models... Using device: cuda:0 Successfully loaded 4 IgFold models. Loaded AntiBERTy model. Traceback (most recent call last): File "igbody_test.py", line 15, in do_renum=True, # Renumber predicted antibody structure (Chothia) File "/data/personal/yankai/LiuChang/IgFold/igfold/IgFoldRunner.py", line 119, in fold truncate_sequences=truncate_sequences, File "/data/personal/yankai/LiuChang/IgFold/igfold/utils/folding.py", line 184, in fold return_attention=True, File "/data/personal/yankai/packages/anaconda3/envs/igfold/lib/python3.7/site-packages/antiberty/AntiBERTyRunner.py", line 81, in embed embeddings[i] = embeddings[i][:, a == 1] RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

[px172]

Move attention_mask to where stores embeddings in AntiBERTyRunner.py

`

    # gather embeddings
    embeddings = outputs.hidden_states
    embeddings = torch.stack(embeddings, dim=1)
    embeddings = list(embeddings.detach().cpu())

    attention_mask = attention_mask.to(embeddings[0].device)

    for i, a in enumerate(attention_mask):
        embeddings[i] = embeddings[i][:, a == 1]

`

[linuxonly801]

Thanks for the replay, it works!

[jeffreyruffolo]

Thanks for pointing this issue out! This should now be fixed in AntiBERTy and IgFold.