skip infer test

hpcaitech · Oct 27, 2023 · 0ef477e · 0ef477e
1 parent c28ee57
commit 0ef477e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/test_infer/test_pipeline_infer.py b/tests/test_infer/test_pipeline_infer.py
@@ -2,12 +2,15 @@
 import torch
 import torch.distributed as dist
 import transformers
+from packaging import version
 
 import colossalai
 from colossalai.inference.pipeline import PPInferEngine
 from colossalai.inference.pipeline.policies import LlamaModelInferPolicy
 from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
 
+CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.5")
+
 
 def data_gen():
  input_ids = torch.tensor([[15496, 11, 616, 3290, 318, 13779, 318, 13779]], dtype=torch.int64)