Skip to content

Commit 6da8f15

Browse files
committed
add back missing post process
1 parent a07d137 commit 6da8f15

File tree

1 file changed

+9
-0
lines changed

1 file changed

+9
-0
lines changed

torchao/dtypes/uintx/uintx_layout.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,15 @@ class UintxLayout(Layout):
228228
dtype: torch.dtype
229229
pack_dim: int = -1
230230

231+
def post_process(
232+
self,
233+
input: torch.Tensor,
234+
scale: torch.Tensor,
235+
zero_point: torch.Tensor,
236+
block_size: Tuple[int, ...],
237+
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
238+
return to_uintx(input, self.dtype, self.pack_dim), scale, zero_point
239+
231240

232241
@register_layout(UintxLayout)
233242
class UintxAQTTensorImpl(PlainAQTTensorImpl):

0 commit comments

Comments
 (0)