fix export

nod-ai · Jan 28, 2025 · 99d3a50 · 99d3a50
1 parent 03df917
commit 99d3a50
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/sharktank/sharktank/layers/linear.py b/sharktank/sharktank/layers/linear.py
@@ -73,6 +73,8 @@ def forward(self, x):
             x = q_input.quantize(x)
             if self.fake_quant:
                 x = x.unpack().dequant()
+            else:
+                x = x.unpack().qs
 
         elif qdq_input is not None:
             x = qdq_input.quantize(x).unpack().dequant()

diff --git a/sharktank/sharktank/ops/qlinear_impls.py b/sharktank/sharktank/ops/qlinear_impls.py
@@ -52,9 +52,7 @@ def qlinear_tensor_scaled(
     if x_layout.qs.dtype.is_floating_point or weight_layout.qs.dtype.is_floating_point:
         if x_layout.qs.dtype == torch.float8_e4m3fnuz:
             # assume quark
-            return matmul(x_layout.qs, weight_layout.qs, transpose_rhs=True).to(
-                torch.float16
-            )
+            return matmul(x_layout.qs, weight_layout.qs, transpose_rhs=True)
         else:
             return NotImplemented