use proprietary dlami for curand

aws · Mar 13, 2024 · 66ce9fc · 66ce9fc
1 parent e5716bc
commit 66ce9fc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/dlc_tests/ec2/test_curand.py b/test/dlc_tests/ec2/test_curand.py
@@ -21,7 +21,7 @@
 @pytest.mark.team("frameworks")
 @pytest.mark.parametrize("ec2_instance_type", CURAND_EC2_SINGLE_GPU_INSTANCE_TYPE, indirect=True)
 @pytest.mark.parametrize(
-    "ec2_instance_ami", [test_utils.UBUNTU_20_BASE_OSS_DLAMI_US_WEST_2], indirect=True
+    "ec2_instance_ami", [test_utils.UBUNTU_20_BASE_PROPRIETARY_DLAMI_US_WEST_2], indirect=True
 )
 def test_curand_gpu(training, ec2_connection, gpu_only, ec2_instance_type):
     if test_utils.is_image_incompatible_with_instance_type(training, ec2_instance_type):