Fixes for failing tests

pytorch · Aug 7, 2024 · 01f0ba1 · 01f0ba1
1 parent 13023d4
commit 01f0ba1
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/test/float8/test_inference_flows.py b/test/float8/test_inference_flows.py
@@ -33,6 +33,7 @@
 random.seed(0)
 torch.manual_seed(0)
 
+is_H100 = torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0)
 is_cuda_8_9 = torch.cuda.is_available() and torch.cuda.get_device_capability() >= (8, 9)
 
 class FeedForward(nn.Module):