Possible 1.5.1 incompatibility fix for elastic inference. See here https://docs.aws.amazon.com/elastic-inference/latest/developerguide/ei-pytorch-using.html

Author: Theiss Heilker

src/sagemaker_pytorch_serving_container/default_pytorch_inference_handler.py

@@ -15,7 +15,7 @@
 import os
 import textwrap
 
-import torch
+import torch, torcheia
 from sagemaker_inference import (
     content_types,
     decoder,
@@ -47,7 +47,9 @@ def default_model_fn(self, model_dir):
                 raise FileNotFoundError("Failed to load model with default model_fn: missing file {}."
                                         .format(DEFAULT_MODEL_FILENAME))
             # Client-framework is CPU only. But model will run in Elastic Inference server with CUDA.
-            return torch.jit.load(model_path, map_location=torch.device('cpu'))
+            model=torch.jit.load(model_path, map_location=torch.device('cpu'))
+            # Attached EIA to model as required for PyTorch 1.5.1
+            return torcheia.jit.attach_eia(model, 0)
         else:
             raise NotImplementedError(textwrap.dedent("""
             Please provide a model_fn implementation.