Address comments

Author: saimidu

test/integration/sagemaker/test_default_inference.py

@@ -102,18 +102,12 @@ def _test_default_inference(
         sagemaker_session=sagemaker_session,
     )
     with timeout_and_delete_endpoint(endpoint_name, sagemaker_session, minutes=30):
-        # Use accelerator type to differentiate EI vs. CPU and GPU. Don't use processor value
-        if accelerator_type is not None:
-            predictor = pytorch.deploy(
-                initial_instance_count=1,
-                instance_type=instance_type,
-                accelerator_type=accelerator_type,
-                endpoint_name=endpoint_name,
-            )
-        else:
-            predictor = pytorch.deploy(
-                initial_instance_count=1, instance_type=instance_type, endpoint_name=endpoint_name
-            )
+        predictor = pytorch.deploy(
+            initial_instance_count=1,
+            instance_type=instance_type,
+            accelerator_type=accelerator_type,
+            endpoint_name=endpoint_name,
+        )
 
         if accelerator_type:
             batch_size = 100

test/resources/mnist/default_model_eia/code/mnist.py

@@ -26,7 +26,7 @@ def predict_fn(input_data, model):
     logger.info('Performing EIA inference with Torch JIT context with input of size {}'.format(input_data.shape))
     # With EI, client instance should be CPU for cost-efficiency.
     # Sub-graphs with unsupported arguments run locally. Server runs with CUDA
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    device = torch.device('cpu')
     model = model.to(device)
     input_data = input_data.to(device)
     with torch.no_grad():