Skip to content

Commit 67635cb

Browse files
committed
skip gather_gemv code check for B200 and fb_code
1 parent 07b1182 commit 67635cb

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

test/test_examples.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
from packaging import version
66
import torch
7+
from torch._environment import is_fbcode
78

89
import helion
910
from helion._testing import DEVICE
@@ -1205,7 +1206,9 @@ def expected(w, idx, x):
12051206
num_stages=1,
12061207
)
12071208

1208-
if is_cuda():
1209+
gpu_name = torch.cuda.get_device_name(0)
1210+
1211+
if is_cuda() and not is_fbcode() and "b200" not in gpu_name.lower():
12091212
self.assertExpectedJournal(code)
12101213

12111214
def test_int4_gemm(self):

0 commit comments

Comments
 (0)