Fix for issue #367

aqlaboratory · Nov 29, 2023 · 2c36cb8 · 2c36cb8
1 parent cdeb8d1
commit 2c36cb8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/openfold/model/structure_module.py b/openfold/model/structure_module.py
@@ -716,7 +716,7 @@ def forward(
         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))