@@ -116,7 +116,7 @@ def test_pytorch_train_mnist_cpu_deep_canary(pytorch_training, ec2_connection, c
116
116
execute_ec2_training_test (ec2_connection , pytorch_training , PT_MNIST_CMD )
117
117
118
118
119
- @pytest .mark .parametrize ("ec2_instance_ami" , [test_utils .UL20_PT_NEURON_US_WEST_2 ], indirect = True )
119
+ @pytest .mark .parametrize ("ec2_instance_ami" , [test_utils .UL20_BASE_NEURON_US_WEST_2 ], indirect = True )
120
120
@pytest .mark .parametrize ("ec2_instance_type" , PT_EC2_NEURON_TRN1_INSTANCE_TYPE , indirect = True )
121
121
@pytest .mark .integration ("pytorch_neuronx_sanity_test" )
122
122
@pytest .mark .neuronx_test
@@ -126,7 +126,7 @@ def test_pytorch_allreduce_neuronx(pytorch_training_neuronx, ec2_connection):
126
126
execute_ec2_training_test (ec2_connection , pytorch_training_neuronx , PT_NEURON_ALLREDUCE_CMD )
127
127
128
128
129
- @pytest .mark .parametrize ("ec2_instance_ami" , [test_utils .UL20_PT_NEURON_US_WEST_2 ], indirect = True )
129
+ @pytest .mark .parametrize ("ec2_instance_ami" , [test_utils .UL20_BASE_NEURON_US_WEST_2 ], indirect = True )
130
130
@pytest .mark .parametrize ("ec2_instance_type" , PT_EC2_NEURON_TRN1_INSTANCE_TYPE , indirect = True )
131
131
@pytest .mark .integration ("pytorch_neuronx_sanity_test" )
132
132
@pytest .mark .neuronx_test
@@ -136,7 +136,7 @@ def test_pytorch_train_mlp_neuronx(pytorch_training_neuronx, ec2_connection):
136
136
execute_ec2_training_test (ec2_connection , pytorch_training_neuronx , PT_NEURON_MLP_CMD )
137
137
138
138
139
- @pytest .mark .parametrize ("ec2_instance_ami" , [test_utils .UL20_PT_NEURON_US_WEST_2 ], indirect = True )
139
+ @pytest .mark .parametrize ("ec2_instance_ami" , [test_utils .UL20_BASE_NEURON_US_WEST_2 ], indirect = True )
140
140
@pytest .mark .parametrize ("ec2_instance_type" , PT_EC2_NEURON_INF2_INSTANCE_TYPE , indirect = True )
141
141
@pytest .mark .integration ("pytorch_neuronx_sanity_test" )
142
142
@pytest .mark .neuronx_test
@@ -146,7 +146,7 @@ def test_pytorch_allreduce_neuronx_inf2(pytorch_training_neuronx, ec2_connection
146
146
execute_ec2_training_test (ec2_connection , pytorch_training_neuronx , PT_NEURON_ALLREDUCE_CMD )
147
147
148
148
149
- @pytest .mark .parametrize ("ec2_instance_ami" , [test_utils .UL20_PT_NEURON_US_WEST_2 ], indirect = True )
149
+ @pytest .mark .parametrize ("ec2_instance_ami" , [test_utils .UL20_BASE_NEURON_US_WEST_2 ], indirect = True )
150
150
@pytest .mark .parametrize ("ec2_instance_type" , PT_EC2_NEURON_INF2_INSTANCE_TYPE , indirect = True )
151
151
@pytest .mark .integration ("pytorch_neuronx_sanity_test" )
152
152
@pytest .mark .neuronx_test
0 commit comments