Multimodel endpoint support with sagemaker

Hi, I'm trying to deploy multiple models with multimodel mode within the same endpoint using torchserve. However, I get the following error when I try to deploy it:

Error hosting endpoint torchserve-endpoint-2021-04-08-19-07-59: Failed. Reason: The primary container for production variant AllTraffic did not pass the ping health check. Please check CloudWatch logs for this endpoint..

When I checked the cloudwatch logs, I could see the following: ACCESS_LOG - /10.32.0.2:48286 "GET /models HTTP/1.1" 404 2

The following is my configuration:

1. Docker file

FROM ubuntu:18.04

ENV PYTHONUNBUFFERED TRUE
LABEL com.amazonaws.sagemaker.capabilities.multi-models=true

RUN apt-get update && \
    DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
    fakeroot \
    ca-certificates \
    dpkg-dev \
    g++ \
    python3-dev \
    openjdk-11-jdk \
    curl \
    vim \
    && rm -rf /var/lib/apt/lists/* \
    && cd /tmp \
    && curl -O https://bootstrap.pypa.io/get-pip.py \
    && python3 get-pip.py

RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1
RUN update-alternatives --install /usr/local/bin/pip pip /usr/local/bin/pip3 1

RUN pip install --no-cache-dir psutil \
                --no-cache-dir torch \
                --no-cache-dir torchvision
                
ADD serve serve
RUN pip install ../serve/

RUN pip install requests


COPY dockerd-entrypoint.sh /usr/local/bin/dockerd-entrypoint.sh
RUN chmod +x /usr/local/bin/dockerd-entrypoint.sh

RUN mkdir -p /opt/ml/model 
RUN mkdir -p /home/model-server/ && mkdir -p /home/model-server/tmp
COPY config.properties /home/model-server/config.properties

WORKDIR /home/model-server
ENV TEMP=/home/model-server/tmp
ENTRYPOINT ["/usr/local/bin/dockerd-entrypoint.sh"]
CMD ["serve"]

2. Config properties

inference_address=http://0.0.0.0:8080
management_address=http://0.0.0.0:8081
number_of_netty_threads=32
job_queue_size=1000
model_store=/opt/ml/model

3. Deploy code

from sagemaker.model import Model
from sagemaker.multidatamodel import MultiDataModel
from sagemaker.predictor import RealTimePredictor

model_data_prefix = "s3://sagemaker-us-east-1-149465543054/sagemaker/image-caption-model"
model_data = "s3://sagemaker-us-east-1-149465543054/sagemaker/image-caption-model/models/caption_model.tar.gz"

sm_model_name = 'torchserve-image-caption-12'
role = 'arn:aws:iam::149465543054:role/ai-rekognition-assume-role'

torchserve_model = Model(model_data = model_data, 
                         image_uri = image,
                         role  = role,
                         predictor_cls=RealTimePredictor,
                         name  = sm_model_name)

multi_model = MultiDataModel(name              = sm_model_name,
                             model_data_prefix = model_data_prefix,
                             model             = torchserve_model)

multi_model.add_model(torchserve_model.model_data)

endpoint_name = 'torchserve-endpoint-' + time.strftime("%Y-%m-%d-%H-%M-%S", time.gmtime())

predictor = multi_model.deploy(instance_type='ml.m4.xlarge',
                               initial_instance_count=1,
                               endpoint_name = endpoint_name)

4. Custom handler initialize function:

class HiggsClassifier:

    def __init__(self):
        self.model = None
        self.device = None
        self.initialized = False
        self.JPEG_CONTENT_TYPE = 'image/jpeg'
        self.JSON_CONTENT_TYPE = 'application/json'
        
    def initialize(self, ctx):
        self.manifest = ctx.manifest
        properties = ctx.system_properties
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model_dir = properties.get("model_dir")
        serialized_file = self.manifest['model']['serializedFile']
        model_pt_path = os.path.join(model_dir, serialized_file)
        #self.model = torch.jit.load(model_pt_path)
        
        
        self.model  = torch.load(model_pt_path, map_location=str(self.device))
        self.decoder = self.model['decoder']
        self.decoder = self.decoder.to(self.device)
        self.decoder.eval()
        self.encoder = self.model['encoder']
        self.encoder = self.encoder.to(self.device)
        self.encoder.eval()
        
        
        word_map_version = os.path.join(model_dir, "WORDMAP.json")
        # Load word map (word2ix)
        with open(word_map_version, 'r') as j:
            self.word_map = json.load(j)
        self.rev_word_map = {v: k for k, v in self.word_map.items()}  # ix2word
        
       
        logger.debug(
            'Model file {0} loaded successfully'.format(model_pt_path))
        self.initialized = True

How do I deploy multiple models to an endpoint using torchserve?