From a4834f02d3c3731fb162c19f1a2d5c50423614e9 Mon Sep 17 00:00:00 2001 From: Li Xing Date: Tue, 16 Jul 2024 14:14:10 +0800 Subject: [PATCH] Update embedders.py change template angle feat for bf16 inference --- fastfold/model/fastnn/embedders.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/fastfold/model/fastnn/embedders.py b/fastfold/model/fastnn/embedders.py index db91f04b..c31888a4 100644 --- a/fastfold/model/fastnn/embedders.py +++ b/fastfold/model/fastnn/embedders.py @@ -185,7 +185,8 @@ def forward(self, template_angle_feat = build_template_angle_feat( single_template_feats, ) - + template_angle_feat = template_angle_feat.to(dtype=z.dtype) + # [*, S_t, N, C_m] a = self.template_angle_embedder(template_angle_feat)