Add guard for cuda-python in new test.

Signed-off-by: Daniel Galvez <[email protected]>
galv · Feb 23, 2024 · e099042 · e099042
1 parent 87d07b8
commit e099042
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/collections/asr/decoding/test_cuda_graph_rnnt_greedy_decoding.py b/tests/collections/asr/decoding/test_cuda_graph_rnnt_greedy_decoding.py
@@ -79,6 +79,8 @@ def test_cuda_graph_rnnt_greedy_decoder(model_name, batch_size, enable_bfloat16)
 
 
 def test_change_devices():
+    skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported()
+
     if torch.cuda.device_count() < 2:
         pytest.skip("Test requires more than 2 GPUs")