Skip to content

Commit

Permalink
Run float8 test only if cuda is available and device capability is gr…
Browse files Browse the repository at this point in the history
…eater than 90
  • Loading branch information
rahul-tuli committed Nov 27, 2024
1 parent 3a6ccc8 commit 74d6498
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion tests/test_utils/test_semi_structured_conversions.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,12 @@


def supported_dtypes():
return [torch.int8, torch.float16, torch.bfloat16, torch.float8_e4m3fn]
dtypes = [torch.int8, torch.float16, torch.bfloat16]
if torch.cuda.is_available():
major, minor = torch.cuda.get_device_capability()
if major > 9 or (major == 9 and minor > 0):
dtypes += [torch.float8_e4m3fn]
return dtypes


def get_random_mat(M, K, dtype):
Expand Down

0 comments on commit 74d6498

Please sign in to comment.