Comment pytest.skip for eia mnist test

Author: Chunxi Wang

buildspec.yml

@@ -89,9 +89,9 @@ phases:
       - execute-command-if-has-matching-changes "$test_cmd" "test/" "src/*.py" "setup.py" "setup.cfg" "buildspec-toolkit.yml" "artifacts/*"
 
     finally:
-      # # shut down remote GPU instance
-      # - cleanup-gpu-instances
-      # - cleanup-key-pairs
+      # shut down remote GPU instance
+      - cleanup-gpu-instances
+      - cleanup-key-pairs
 
       # # remove ECR image
       # - aws ecr batch-delete-image --repository-name $ECR_REPO --region $AWS_DEFAULT_REGION --image-ids imageTag=$GENERIC_TAG

test/integration/sagemaker/test_mnist.py

@@ -34,7 +34,7 @@ def test_mnist_gpu(sagemaker_session, image_uri, instance_type):
     _test_mnist_distributed(sagemaker_session, image_uri, instance_type, model_gpu_tar, mnist_gpu_script)
 
 
-@pytest.mark.skip(reason="Latest EIA version is too old - 1.3.1. Remove this after a new DLC release")
+# @pytest.mark.skip(reason="Latest EIA version is too old - 1.3.1. Remove this after a new DLC release")
 @pytest.mark.eia_test
 def test_mnist_eia(sagemaker_session, image_uri, instance_type, accelerator_type):
     instance_type = instance_type or 'ml.c4.xlarge'

test/resources/mnist/model_eia/mnist.py

@@ -31,6 +31,7 @@ def predict_fn(input_data, model):
     input_data = input_data.to(device)
     model = model.eval() 
     with torch.no_grad():
+        print("current torch version is: ", torch.__version__)
         if torch.__version__ in VERSIONS_USE_NEW_API:
             import torcheia
             # we need to set the profiling executor for EIA