maybe?

chenyuxyz · May 15, 2024 · e89c73c · e89c73c
1 parent 5113c84
commit e89c73c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/extra/models/bert.py b/extra/models/bert.py
@@ -64,7 +64,7 @@ def __call__(self, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
     onehot = counter == masked_positions.unsqueeze(2).expand(*masked_positions.shape, output.shape[1])
     h_masked = onehot @ output
 
-    h_masked = self.lm_norm(self.lm_transform_activation(self.lm_transform(h_masked)))
+    h_masked = self.lm_norm(self.lm_transform_activation(self.lm_transform(h_masked).float()))
     lm_logits = self.lm_output(h_masked) + self.lm_output_bias
 
     return lm_logits, clsf_logits