Не получается импортировать мой класс архитектуры модели Pytorch в мой код класса пользовательского предиктора для развертывания пользовательского конвейера на платформе AI

Question

У меня проблемы с развертыванием простой обученной модели Pytorch на AI Platform Serving.

Вот ошибка:

Creating version (this might take a few minutes)......failed.
ERROR: (gcloud.beta.ai-platform.versions.create) Create Version failed. Bad model detected with error:  "Failed to load model: Unexpected error when loading the model: Can't get attribute 'Net' 
on <module '__main__' from 'prediction_server_beta.py'> (Error code: 0)"

Вот класс пользовательского предсказания:

import os
import pickle 
import numpy as np 
import torch
from torch_model import *
class CustomModelPrediction(object):
    def __init__(self, model):
        self._model = model

    def predict(self, instances):
        input = torch.Tensor(instances)
        predictions = self._model(input)

        return predictions

    @classmethod
    def from_path(cls, model_dir):
        from torch_model import Net
        model = Net()
        model = torch.load(os.path.join(model_dir, 'pickle_saved_model.pkl'))
        #model.eval()
        return cls(model)

Я попытался изменить код и сохранить форматы моделей (.pkl,. pth, .pt), но, похоже, ничего не работает. Я также попытался включить мой класс модели в тот же сценарий .py, что и мой класс пользовательских предикторов, но это тоже не сработало. Пакет для установки в pip содержал все необходимые коды, то есть код модели, а также код пользовательского предиктора. Ценю помощь!