Missing .py file causing attributeError on lambda

[shawn]

I want to build a docker image to add my deep learning model GRU into my lambda function.
Here is the GRUModel.py

import torch
import torch.nn as nn
class GRUModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, layer_dim, output_dim, dropout_prob):
        super(GRUModel, self).__init__()
        # Defining the number of layers and the nodes in each layer
        self.layer_dim = layer_dim
        self.hidden_dim = hidden_dim
        # GRU layers
        self.gru = nn.GRU(
            input_dim, hidden_dim, layer_dim, batch_first=True, dropout=dropout_prob
        )
        # Fully connected layer
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        # Initializing hidden state for first input with zeros
        h0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).requires_grad_()
        # Forward propagation by passing in the input and hidden state into the model
        out, _ = self.gru(x, h0.detach())
        # Reshaping the outputs in the shape of (batch_size, seq_length, hidden_size)
        # so that it can fit into the fully connected layer
        out = out[:, -1, :]
        # Convert the final state to our desired output shape (batch_size, 
         output_dim)
        out = self.fc(out)
        return out

Here is the lambda.py

import torch
import torch.nn as nn
import joblib
from GRUModel import GRUModel
def handler(event, context):
    gruencoder=joblib.load("gruencoder.pkl")
    response = {'statusCode': 200, 'body' : "OK"}
    return response

Here is the Dockerfile

FROM public.ecr.aws/lambda/python:3.7
COPY lambda.py ${LAMBDA_TASK_ROOT}
COPY gruencoder.pkl .
COPY GRUModel.py .
RUN pip3 install joblib --target "${LAMBDA_TASK_ROOT}"
RUN pip3 install torch --target "${LAMBDA_TASK_ROOT}"
CMD ["lambda.handler"]

I run the lambda.py is working on the loacl, but it shows error on the lambda.

{
  "errorMessage": "module '__main__' has no attribute 'GRUModel'",
  "errorType": "AttributeError",
  "stackTrace": [
    "  File \"/var/task/lambda.py\", line 22, in handler\n    gruencoder=joblib.load(\"gruencoder.pkl\")\n",
    "  File \"/var/task/joblib/numpy_pickle.py\", line 587, in load\n    obj = _unpickle(fobj, filename, mmap_mode)\n",
    "  File \"/var/task/joblib/numpy_pickle.py\", line 506, in _unpickle\n    obj = unpickler.load()\n",
    "  File \"/var/lang/lib/python3.7/pickle.py\", line 1088, in load\n    dispatch[key[0]](self)\n",
    "  File \"/var/lang/lib/python3.7/pickle.py\", line 1376, in load_global\n    klass = self.find_class(module, name)\n",
    "  File \"/var/lang/lib/python3.7/pickle.py\", line 1430, in find_class\n    return getattr(sys.modules[module], name)\n"
  ]
}