From 2676476013a7a57eff0bfe5030bde60cc6216a69 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sun, 31 May 2026 21:51:59 +0000
Subject: [PATCH] fix mHC pre_block bmm dtype mismatch: A is FP32, X is BF16

---
 single_shot_inference.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/single_shot_inference.py b/single_shot_inference.py
index 6205550b..0ae5dfc3 100644
--- 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):