Could not export model using torch.jit.trace

[Songweiping]

trafficstars

Hi authors，

Thanks for you great work! I trained a model based on my own dataset, but encountered a problem during exporting the model. Could not export Python function call 'MambaInnerFn'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__

It looks like MambaInnerFn is not supported by torch.jit.trace, do you have any plan or suggestion to fix this?

[tridao]

Yeah I don't think torch.jit.trace is supported by our custom function.

[Songweiping]

Thanks for your prompt reply. I'm wondering is it possible to make it compatible with torch.jit.trace for deployment need?

[tridao]

I'm not familiar with torch.jit.trace, but we welcome community's contribution.