Fix for issue #367

2c36cb8b · Christina Floristean · cdeb8d1b · 2c36cb8b
Commit 2c36cb8b authored Nov 29, 2023 by Christina Floristean
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

openfold/model/structure_module.py openfold/model/structure_module.py +1 -1

No files found.
--- a/openfold/model/structure_module.py
+++ b/openfold/model/structure_module.py
@@ -716,7 +716,7 @@ class InvariantPointAttentionMultimer(nn.Module):
        o_pt_norm = o_pt.norm(epsilon=1e-8)

        if (_offload_inference):
-            z[0] = z[0].to(o_pt.device)
+            z[0] = z[0].to(o_pt.x.device)

        o_pair = torch.einsum('...ijh, ...ijc->...ihc', a, z[0].to(dtype=a.dtype))