fix mHC pre_block bmm dtype mismatch: A is FP32, X is BF16

2026-05-31 21:51:59 +00:00
parent eb08cd06d1
commit 2676476013
1 changed files with 1 additions and 1 deletions
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -157,7 +157,7 @@ class mHCBlock:
        A = torch.sigmoid(pre_t) + HC_EPS
        C = 2.0 * torch.sigmoid(post_t)
        B = sinkhorn_knopp(comb_t.reshape(T, n, n), t_max=self.t_max)
-        x_in = torch.bmm(A.unsqueeze(1), X).squeeze(1).bfloat16()
+        x_in = torch.bmm(A.unsqueeze(1), X.float()).squeeze(1).bfloat16()
        return x_in, {'B': B, 'C': C}

    def post_block(self, X, F_out, ctx):