title: BERT model loading not working with pytorch 1.3.1-eia container
body: <p>I have a custom python file for inference in which I have implemented the functions <code>model_fn</code>, <code>input_fn</code>, <code>predict_fn</code> and <code>output_fn</code>.  I have saved the model as a torchscript using <code>torch.jit.trace</code>, <code>torch.jit.save</code> and loading it using <code>torch.jit.load</code>. The <code>model_fn</code> implementation is as follows:</p>
  <pre><code>import torch
  import os
  import logging
  
  logger = logging.getLogger()
  is_ei = os.getenv(&quot;SAGEMAKER_INFERENCE_ACCELERATOR_PRESENT&quot;) == &quot;true&quot;
  logger.warn(f&quot;Elastic Inference enabled: {is_ei}&quot;)
  
  def model_fn(model_dir):
      model_path = os.path.join(model_dir, &quot;model_best.pt&quot;)
      try:
          loaded_model = torch.jit.load(model_path, map_location=torch.device('cpu'))
          loaded_model.eval()
          return loaded_model
      except Exception as e:
          logger.exception(f&quot;Exception in model fn {e}&quot;)
          return None
  </code></pre>
  <p>This implementation works perfectly for the container with pytorch 1.5. But for container with torch 1.3.1 it exits abruptly when loading the pretrained model without any logs. The only line I see in the logs is</p>
  <pre><code>algo-1-nvqf7_1  | 2020-11-30 07:17:15,392 [WARN ] W-9000-model com.amazonaws.ml.mms.wlm.BatchAggregator - Load model failed: model, error: Worker died.
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,393 [INFO ] W-9000-model com.amazonaws.ml.mms.wlm.WorkerThread - Retry worker: 9000-44f1cd64 in 1 seconds.
  </code></pre>
  <p>The worker dies and tries to restart, and the process repeats till I stop the container.</p>
  <p>The model I am using is trained with pytorch 1.5. But since EI support is only supported till 1.3.1, I am using this container.</p>
  <p>Things I have tried:</p>
  <ol>
  <li>The same code with same model works outside the container with pytorch version 1.3.1. So, I don't think pytorch version compatibility is the issue.</li>
  <li>Tried using <code>debug</code> and <code>notset</code> levels for logs. Didn't get any more info as to why model loading fails</li>
  <li>Tried loading the original model instead of the traced one. Again this works in 1.5 but not in 1.3.1. Fails at the same point, while loading the BERT pretrained model.</li>
  <li>Tried this setup on sagemaker notebook instance with gpu accelerator and sagemaker <code>PytorchModel's deploy()</code> function with <code>framework_version</code> as 1.3.1. Also tried it using the 1.3.1 container without <code>eia</code>. Has same behaviour everywhere.</li>
  </ol>
  <p>Am I doing something wrong or missing something crucial from the documentation? Any help would be much appreciated.</p>
  <p>**Logs for container with torch 1.3.1-eia **</p>
  <pre><code>algo-1-nvqf7_1  | 2020-11-30 07:17:14,333 [INFO ] main com.amazonaws.ml.mms.ModelServer - 
  algo-1-nvqf7_1  | MMS Home: /opt/conda/lib/python3.6/site-packages
  algo-1-nvqf7_1  | Current directory: /
  algo-1-nvqf7_1  | Temp directory: /home/model-server/tmp
  algo-1-nvqf7_1  | Number of GPUs: 0
  algo-1-nvqf7_1  | Number of CPUs: 8
  algo-1-nvqf7_1  | Max heap size: 6972 M
  algo-1-nvqf7_1  | Python executable: /opt/conda/bin/python
  algo-1-nvqf7_1  | Config file: /etc/sagemaker-mms.properties
  algo-1-nvqf7_1  | Inference address: http://0.0.0.0:8080
  algo-1-nvqf7_1  | Management address: http://0.0.0.0:8080
  algo-1-nvqf7_1  | Model Store: /.sagemaker/mms/models
  algo-1-nvqf7_1  | Initial Models: ALL
  algo-1-nvqf7_1  | Log dir: /logs
  algo-1-nvqf7_1  | Metrics dir: /logs
  algo-1-nvqf7_1  | Netty threads: 0
  algo-1-nvqf7_1  | Netty client threads: 0
  algo-1-nvqf7_1  | Default workers per model: 1
  algo-1-nvqf7_1  | Blacklist Regex: N/A
  algo-1-nvqf7_1  | Maximum Response Size: 6553500
  algo-1-nvqf7_1  | Maximum Request Size: 6553500
  algo-1-nvqf7_1  | Preload model: false
  algo-1-nvqf7_1  | Prefer direct buffer: false
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,391 [WARN ] W-9000-model com.amazonaws.ml.mms.wlm.WorkerLifeCycle - attachIOStreams() threadName=W-9000-model
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,481 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - model_service_worker started with args: --sock-type unix --sock-name /home/model-server/tmp/.mms.sock.9000 --handler sagemaker_pytorch_serving_container.handler_service --model-path /.sagemaker/mms/models/model --model-name model --preload-model false --tmp-dir /home/model-server/tmp
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,482 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Listening on port: /home/model-server/tmp/.mms.sock.9000
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,482 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - [PID] 51
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,482 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - MMS worker started.
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,483 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Python runtime: 3.6.6
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,483 [INFO ] main com.amazonaws.ml.mms.wlm.ModelManager - Model model loaded.
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,487 [INFO ] main com.amazonaws.ml.mms.ModelServer - Initialize Inference server with: EpollServerSocketChannel.
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,496 [INFO ] W-9000-model com.amazonaws.ml.mms.wlm.WorkerThread - Connecting to: /home/model-server/tmp/.mms.sock.9000
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,544 [INFO ] main com.amazonaws.ml.mms.ModelServer - Inference API bind to: http://0.0.0.0:8080
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,545 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Connection accepted: /home/model-server/tmp/.mms.sock.9000.
  algo-1-nvqf7_1  | Model server started.
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,547 [WARN ] pool-2-thread-1 com.amazonaws.ml.mms.metrics.MetricCollector - worker pid is not available yet.
  algo-1-nvqf7_1  | 2020-11-30 07:17:14,962 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - PyTorch version 1.3.1 available.
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,314 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Lock 140580224398952 acquired on /root/.cache/torch/transformers/26bc1ad6c0ac742e9b52263248f6d0f00068293b33709fae12320c0e35ccfbbb.542ce4285a40d23a559526243235df47c5f75c197f04f37d1a0c124c32c9a084.lock
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,315 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-vocab.txt not found in cache or force_download set to True, downloading to /root/.cache/torch/transformers/tmpcln39mxo
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,344 [WARN ] W-9000-model-stderr com.amazonaws.ml.mms.wlm.WorkerLifeCycle - 
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,349 [WARN ] W-9000-model-stderr com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Downloading:   0%|          | 0.00/232k [00:00&lt;?, ?B/s]
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,349 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - storing https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-vocab.txt in cache at /root/.cache/torch/transformers/26bc1ad6c0ac742e9b52263248f6d0f00068293b33709fae12320c0e35ccfbbb.542ce4285a40d23a559526243235df47c5f75c197f04f37d1a0c124c32c9a084
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,349 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - creating metadata file for /root/.cache/torch/transformers/26bc1ad6c0ac742e9b52263248f6d0f00068293b33709fae12320c0e35ccfbbb.542ce4285a40d23a559526243235df47c5f75c197f04f37d1a0c124c32c9a084
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,349 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Lock 140580224398952 released on /root/.cache/torch/transformers/26bc1ad6c0ac742e9b52263248f6d0f00068293b33709fae12320c0e35ccfbbb.542ce4285a40d23a559526243235df47c5f75c197f04f37d1a0c124c32c9a084.lock
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,350 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - loading file https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-vocab.txt from cache at /root/.cache/torch/transformers/26bc1ad6c0ac742e9b52263248f6d0f00068293b33709fae12320c0e35ccfbbb.542ce4285a40d23a559526243235df47c5f75c197f04f37d1a0c124c32c9a084
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,378 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Created tokenizer
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,378 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Elastic Inference enabled: True
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,378 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - inside model fn
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,379 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - /.sagemaker/mms/models/model
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,379 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - /.sagemaker/mms/models/model/model.pt
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,379 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - ['model.pt', 'model.tar.gz', 'code', 'model_tn_best.pth', 'MAR-INF']
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,379 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Loading torch script
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,392 [INFO ] epollEventLoopGroup-4-1 com.amazonaws.ml.mms.wlm.WorkerThread - 9000-44f1cd64 Worker disconnected. WORKER_STARTED
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,392 [WARN ] W-9000-model com.amazonaws.ml.mms.wlm.BatchAggregator - Load model failed: model, error: Worker died.
  algo-1-nvqf7_1  | 2020-11-30 07:17:15,393 [INFO ] W-9000-model com.amazonaws.ml.mms.wlm.WorkerThread - Retry worker: 9000-44f1cd64 in 1 seconds.
  algo-1-nvqf7_1  | 2020-11-30 07:17:16,065 [INFO ] W-9000-model ACCESS_LOG - /172.18.0.1:45110 &quot;GET /ping HTTP/1.1&quot; 200 8
  </code></pre>