Trouble Exporting to Torchscript

[will-rice]

I'm having trouble exporting ScatteringTorch1D to torchscript. The module traces fine, but then errors on saving the traced module.

The specific error is this:

RuntimeError:
Could not export Python function call 'ModulusStable'. 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: