Change Predictor Class on Model object

Author: saimidu

test/integration/sagemaker/test_mnist_default_inference.py

@@ -17,7 +17,8 @@
 import pytest
 import requests
 import sagemaker
-from sagemaker.pytorch import PyTorchModel
+from sagemaker.predictor import RealTimePredictor
+from sagemaker.pytorch import PyTorchModel, PyTorchPredictor
 
 from integration import (
     model_cpu_tar,
@@ -85,6 +86,7 @@ def _test_default_inference(
     pytorch = PyTorchModel(
         model_data=model_data,
         role="SageMakerRole",
+        predictor_cls=RealTimePredictor if not accelerator_type else PyTorchPredictor,
         entry_point=mnist_script,
         image=image_uri,
         sagemaker_session=sagemaker_session,

test/resources/resnet18/default_model/code/resnet18.py

@@ -15,7 +15,7 @@ def transform_fn(model, payload, request_content_type, response_content_type):
 
     logger.info("Invoking user-defined transform function")
 
-    if request_content_type != "application/octet-stream":
+    if request_content_type and request_content_type != "application/octet-stream":
         raise RuntimeError(
             "Content type must be application/octet-stream. Provided: {0}".format(
                 request_content_type