How to deploy torchscript module when original data structure class is not avaliable

Take detectron2 as an example. I have transformed a model from detectron2 with torch.jit.script to torchscript, the input of which is annotated with detectron2.structures.box. Because there is no detectron2 python package in the deployment environment of the torchscript module , I define another input data class x.box which has totally the same codes as detectron2.structures.box. But when this new input class is sent to the torchscript module, an error happened:
RuntimeError: classType INTERNAL ASSERT FAILED at /pytorch/torch/csrc/jit/python/pybind_utils.h:742, please report a bug to PyTorch.

Is this mean that I must send the same input data class(detectron2.structures.box) to the model? Or, is there any ways to make it work?

Thanks a lot:)

I have figured out the reason. It is because the data class of the output of torchscript is not supported in the deployment environment.

So is there any way to register a new data class which consist of the same codes of the original one in torchscript in the deployment environment and construct links between them?

Are you able to post a repro for the internal assert failure?

It sounds like you’re asking for structural typing instead of nominal typing. TorchScript is nominally typed, so we don’t have a blessed way to do exactly what you’re asking. (It does look like some version of structural typing is supported, but that is likely an accident).

detectron2.structures.box should be included with the serialized code of your TorchScript model, so you should be able to reconstruct it inside the JIT (though the API to do this may not exist/be pretty). Once we fix the above bug we can probably figure out some hack to make this still work.

The codes are still being modified, but I can illustrate this problem with some sample codes:

In the development environment where torchscript is exported, we have:

from detectron2.structures import Box

class MyModule(nn.Module) -> Box:
    some_of_codes

then we export torchscript:

model = MyModule()
torch.jit.script(model).save('scripted_model.pt')

In the deployment environment, we have:

model = torch.jit.load('scripted_model.pt')
input = some_data_of_model
output = model(input)

Since the data type ‘detectron2.structures.Box’ is not available here, the RuntimeError mentioned above will appear.

I can see that the ‘detectron2.structures.Box’ class is available in the serialized code of my TorchScript, but it seems that I can not use it to construct an instance which can receive the result from the output of torchscript.