From ea7c14f8ef64924f2d0ff80df3cdabf2c7299848 Mon Sep 17 00:00:00 2001 From: Titus von Koeller Date: Mon, 1 Aug 2022 09:32:47 -0700 Subject: reran black with linelength 80 for greater readability --- quicktest.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) (limited to 'quicktest.py') diff --git a/quicktest.py b/quicktest.py index 29d045d..0fcda64 100644 --- a/quicktest.py +++ b/quicktest.py @@ -14,23 +14,31 @@ def test_igemmlt(dim1, dim2, dim3, dim4, dims, ldb): torch.int8 ) elif dims == 3: - A = torch.randint(-128, 127, size=(dim1, dim2, dim3), device="cuda").to( - torch.int8 - ) - B = torch.randint(-128, 127, size=(dim4, dim3), device="cuda").to(torch.int8) + A = torch.randint( + -128, 127, size=(dim1, dim2, dim3), device="cuda" + ).to(torch.int8) + B = torch.randint(-128, 127, size=(dim4, dim3), device="cuda").to( + torch.int8 + ) C1 = torch.matmul(A.float(), B.t().float()) A2, SA = F.transform(A, "col32") B2, SB = F.transform(B, "colx") if dims == 2: C2, SC = F.transform( - torch.zeros(A.shape[0], B.shape[0], dtype=torch.int32, device="cuda"), + torch.zeros( + A.shape[0], B.shape[0], dtype=torch.int32, device="cuda" + ), "col32", ) else: C2, SC = F.transform( torch.zeros( - A.shape[0], A.shape[1], B.shape[0], dtype=torch.int32, device="cuda" + A.shape[0], + A.shape[1], + B.shape[0], + dtype=torch.int32, + device="cuda", ), "col32", ) -- cgit v1.2.3