Error during exporting the model

[Songweiping]

Hi authors,

Thanks for the great work! I use flash-fft-conv in my project and it works quite well. During tracing the model, I encounter a RuntimeError:
Could not export Python function call 'FlashFFTConvFunc'. 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__

Do you know how to fix this issue? Thanks!