Skip to content

GitLab

"...python/git@developer.sourcefind.cn:OpenDAS/tilelang.git" did not exist on "c30df2a1c58bc6296e2a6027b4ebacf9f1b82202"
Commit 5a7db20f authored by zhuww's avatar zhuww
Browse files

fix some problems about struture_module

parent 5e4d790d
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 2 deletions
fastfold/model/nn/structure_module.py
View file @ 5a7db20f
...@@ -518,7 +518,7 @@ class InvariantPointAttention(nn.Module): ...@@ -518,7 +518,7 @@ class InvariantPointAttention(nn.Module):
if(_offload_inference): if(_offload_inference):
z[0] = z[0].to(o_pt.device) z[0] = z[0].to(o_pt.device)
# [*, N_res, H, C_z] # [*, N_res, H, C_z]
o_pair = torch.matmul(a.transpose(-2, -3), z.to(dtype=a.dtype)) o_pair = torch.matmul(a.transpose(-2, -3), z[0].to(dtype=a.dtype))
del a del a
torch.cuda.empty_cache() torch.cuda.empty_cache()
...@@ -774,7 +774,6 @@ class StructureModule(nn.Module): ...@@ -774,7 +774,6 @@ class StructureModule(nn.Module):
s = self.layer_norm_s(s) s = self.layer_norm_s(s)
# [*, N, N, C_z] # [*, N, N, C_z]
z = self.layer_norm_z(z)
# inplace z # inplace z
evoformer_output_dict["pair"] = evoformer_output_dict["pair"].contiguous() evoformer_output_dict["pair"] = evoformer_output_dict["pair"].contiguous()
torch.cuda.empty_cache() torch.cuda.empty_cache()
......
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