There isn’t a way to extract an nn.Module from a compiled ScriptModule. Is it possible for you to instead export your original module instead of a ScriptModule?
For some background, torch.onnx.export will use torch.jit.trace to get an exportable graph from an nn.Module. The ONNX exporter does not support all the features of TorchScript (e.g. if you used torch.jit.script to compile your model, it may not be possible to export that compiled module to ONNX), but relying on torch.jit.trace enforces that only supported features are used.