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):