Error Handling with the Torchscript serialized code

Hello folks -

I want some advice on error handling with torchscript code

  1. Sample Code snippet which validates the user input in forward() and throws ValueError for invalid inputs. Serialize this code using Torchscript
import torch
from torch import nn

class Foo(nn.Module):
  def forward(self, user_input_length: int):
       if user_input_length > 512:
          raise ValueError(f" Invalid user input. Max length allowed 512")
foo_ts = torch.jit.script(Foo())  
torch.jit.save(foo_ts, "foo_ts.pt")
  1. Sharing the torchscript file with the client who loads it in their environment.
foo_ts = torch.jit.load("foo_ts.pt")
foo_ts(198) # Valid Input
foo_ts(513) # Invalid Input - Let the user know
  1. RuntimeError seen by the client when they pass an invalid user input
Traceback (most recent call last):
  File "ts_exception.py", line 11, in <module>
    foo_ts(513)
  File "/home/anjch/.local/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
torch.jit.Error: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "ts_exception.py", line 8, in forward
    def forward(self, user_input_length: int):
         if user_input_length > 512:
            raise ValueError(f" Invalid user input. Max length allowed 512")
            ~~~~~~~~~~~~~~~~~~ <--- HERE
RuntimeError:  Invalid user input. Max length allowed 512

In the current behavior, even if the actual code throws ValueError, client side always receives a RuntimeError. Due to this behavior, the client will not be able to programmatically distinguish between an actual invalid input error v/s other Runtime error (unless they parse the string error message)

Are there any suggestions on how to handle it gracefully where the client can distinguish between input errors vs other errors?