diff --git a/python/ray/dag/tests/experimental/test_torch_tensor_dag.py b/python/ray/dag/tests/experimental/test_torch_tensor_dag.py index 871906fc1f1a3..5fb188b830906 100644 --- a/python/ray/dag/tests/experimental/test_torch_tensor_dag.py +++ b/python/ray/dag/tests/experimental/test_torch_tensor_dag.py @@ -1125,6 +1125,10 @@ def test_torch_tensor_nccl_all_reduce_wrong_shape(ray_start_regular): # exception, such as when the task returns torch.Tensors of the wrong # shape or dtype. Check that we can no longer submit to the DAG. ref = compiled_dag.execute([((20,), dtype, 1) for _ in workers]) + try: + ray.get(ref) + except Exception: + pass with pytest.raises(RayChannelError): ref = compiled_dag.execute([((20,), dtype, 1) for _ in workers])