From 5f41c1eeac85237c873125d7d0c9ac00a262898c Mon Sep 17 00:00:00 2001 From: "Jiang, Yanbing" Date: Thu, 18 Jul 2024 23:12:59 -0700 Subject: [PATCH] Fix CI --- test/test_ops.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/test_ops.py b/test/test_ops.py index 1420e1df8..26f1b8414 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -271,6 +271,8 @@ def test_dequantize_tensor_core_tiled_layout_op(shape, inner_k_tiles, group_size device = "cuda" q = torch.randint(0, 16, shape, dtype=torch.int, device=device) + if TORCH_VERSION_AFTER_2_5: + q = (q[::, ::2] << 4 | q[::, 1::2]).to(torch.uint8) packed_w = torch._convert_weight_to_int4pack(q, inner_k_tiles) q_groups = k // group_size scales = torch.randn(n, q_groups, dtype=torch.bfloat16, device=device)