torch.jit error in get_embedding.py

[wideglide]

There seems to be an error with annotations when using the command/inference/get_embedding.py script.
See the error message:

Traceback (most recent call last):
  File "command/inference/get_embedding.py", line 53, in <module>
    emb0_rep = loaded(sample0_emb, features_only=True, classification_head_name='similarity')[0]['features']
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
torch.jit.Error: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/fairseq/modules/trex_encoder.py", line 252, in forward
    else:
      pass
    ops.prim.RaiseException(_44)
    ~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    return _45
class ByteCombineCNN(Module):

Traceback of TorchScript, original code (most recent call last):
  File "/home/user/trex/fairseq/modules/trex_encoder.py", line 166, in forward

        if self.layernorm_embedding is not None:
            x = self.layernorm_embedding(x)
                ~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        x = self.dropout_module(x)
        if self.quant_noise is not None:
RuntimeError: This Python function is annotated to be ignored and cannot be run

[peikexin9]

Hi @wideglide Thanks for raising the issue. It runs with no problem on my end. What torch version you are using?

[wideglide]

torch version 1.10.1 It works if I comment out the jit code, but not when it tries to use the torch script.

[peikexin9]

hmm this is weird. I am also using pytorch 1.10.1 py3.9_cuda11.3_cudnn8.2.0_0, but it runs no problem with torchscript. Is it the latest code?

[wideglide]

yes, It seemed to be something botched with the virtualenv. I deleted the venv, started over, and now it works? Feel free to close this issue, sorry for the hassle!

[peikexin9]

Not at all @wideglide :-)