fix some problems about struture_module

fastfold/model/nn/structure_module.py

@@ -518,7 +518,7 @@ class InvariantPointAttention(nn.Module):
         if(_offload_inference):
             z[0] = z[0].to(o_pt.device)
         # [*, 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
 
         torch.cuda.empty_cache()
@@ -774,7 +774,6 @@ class StructureModule(nn.Module):
         s = self.layer_norm_s(s)
 
         # [*, N, N, C_z]
-        z = self.layer_norm_z(z)
         # inplace z
         evoformer_output_dict["pair"] = evoformer_output_dict["pair"].contiguous()
         torch.cuda.empty_cache()