"...git@developer.sourcefind.cn:OpenDAS/dynamo.git" did not exist on "e5db9e8634e242ec7965c522ee0bde3034c4a8a0"
Commit 2c36cb8b authored by Christina Floristean's avatar Christina Floristean
Browse files

Fix for issue #367

parent cdeb8d1b
...@@ -716,7 +716,7 @@ class InvariantPointAttentionMultimer(nn.Module): ...@@ -716,7 +716,7 @@ class InvariantPointAttentionMultimer(nn.Module):
o_pt_norm = o_pt.norm(epsilon=1e-8) o_pt_norm = o_pt.norm(epsilon=1e-8)
if (_offload_inference): 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)) o_pair = torch.einsum('...ijh, ...ijc->...ihc', a, z[0].to(dtype=a.dtype))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment