TorchScript Inference fails with torch.jit.trace for seq2seq PyTorch model

harinarayanan_k: <p><strong>Issue:</strong> Inference fails with torch.jit.trace on seq2seq PyTorch model using Bi-LSTM and LSTM</p>
<p><strong>Model architecture</strong>: <a href="https://github.com/LeePleased/StackPropagation-SLU/blob/master/utils/module.py" class="inline-onebox" rel="noopener nofollow ugc">StackPropagation-SLU/module.py at master · LeePleased/StackPropagation-SLU · GitHub</a></p>
<p><strong>Model Input:</strong> Tuple of Input Tensors (Tensor1, Tensor2)  where,<br>
Tensor1 → Words of an Input sentences mapped to their Ids from dictionary<br>
Tensor2 → is the length of original input text.</p>
<p><strong>//Model Instance</strong><br>
import torch<br>
from utils.module import ModelManager<br>
path = ‘save/model/model.pt’<br>
kwargs, state = torch.load(path, map_location=torch.device(‘cpu’))<br>
model = ModelManager(**kwargs)<br>
model.eval()<br>
model.load_state_dict(state)</p>
<p><strong>//Model Input Sample</strong><br>
text = “bring some clay on your return please”<br>
length = 30<br>
X, L = tokenize(text)<br>
X = torch.tensor(X, dtype=torch.long).unsqueeze(0)<br>
L = torch.tensor(L, dtype=torch.long).unsqueeze(0)<br>
inputs = (X,L)</p>
<p><strong>//Torch JIT Trace</strong><br>
traced_script_module = torch.jit.trace(model, inputs)<br>
torch.jit.save(traced_script_module, ‘jitTest.pt’)<br>
res = traced_script_module(X,L) → Throws error during inference for different values of “L” (original text length)</p>
<p><strong>//Error stack trace</strong></p>
<ol>
<li>
<p>When Input length is more than the sample input given during tracing the model:<br>
Input: "text = “bring some clay on your return to home please”<br>
//error<br>
<strong>RuntimeError: The following operation failed in the TorchScript interpreter</strong><br>
<strong>Traceback of TorchScript (most recent call last):</strong><br>
<strong>module.py(305): forward</strong><br>
<strong>Anaconda\lib\site-packages\torch\nn\modules\module.py(1118): _slow_forward</strong><br>
<strong>Anaconda\lib\site-packages\torch\nn\modules\module.py(1130): _call_impl</strong><br>
<strong>module.py(113): forward</strong><br>
<strong>RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 9 but got size 7 for tensor number 1 in the list.</strong></p>
</li>
<li>
<p>When Input length is less than the sample input given during tracing the model:<br>
Input: "text = “bring some clay on your return”<br>
//error<br>
<strong>RuntimeError: The following operation failed in the TorchScript interpreter.</strong><br>
<strong>Traceback of TorchScript (most recent call last):</strong><br>
<strong>module.py(347): forward</strong><br>
<strong>Anaconda\lib\site-packages\torch\nn\modules\module.py(1118): _slow_forward</strong><br>
<strong>Anaconda\lib\site-packages\torch\nn\modules\module.py(1130): _call_impl</strong><br>
<strong>module.py(102): forward</strong><br>
<strong>RuntimeError: index 6 is out of bounds for dimension 0 with size 6</strong></p>
</li>
</ol>
<p><strong>Observation:</strong></p>
<ul>
<li>The issue comes form the decoder part of the model.</li>
<li>Not able to find any function or method calls in model definition which makes the Input Tensor shape dimension constant</li>
<li>The issue arises due to freezing of the dimension of input tensor during tracing which ideally should not happen.</li>
<li>torch.jit.script is not possible here as the model involves inheritance etc and is too complex</li>
</ul>
<p>Kindly request to help if it is a bug from the torchscript tracing module (or) any method or function calls in model architecture is missed which is making the input dimension of tensor constant thereby giving no chance for dynamism. Thanks in advance.</p>