Is there a way to export using torch.jit.script ?

[jprobichaud]

Thanks for this great repository!

I was wondering if there a way to export using torchscript? I tried a simple approach with torch.jit.script(model), but I get:

RuntimeError: 
Module 'RWKV_Block' has no attribute 'att_group_norm' :
  File "/data/workspaces/jp/LLMs/RWKV_Pytorch/src/model.py", line 229
        # 展平x并应用组归一化和门控
        if self.onnx_opset >= 18:
            x = self.att_group_norm(x.flatten(start_dim=1)) * g
                ~~~~~~~~~~~~~~~~~~~ <--- HERE
        else:
            x = x.flatten(start_dim=1) 
'RWKV_Block.time_mixing' is being compiled since it was called from 'RWKV_Block.forward'
  File "/data/workspaces/jp/LLMs/RWKV_Pytorch/src/model.py", line 319
        """
        if self.onnx_opset >= 17:
            x = x + self.time_mixing(self.ln1(x), state, i)
                    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
            x = x + self.channel_mixing(self.ln2(x), state, i)
        else:

[uniartisan]

Hey, we are moving to https://github.com/TorchRWKV/TorchRWKV6

To solve your problem, you can simply checkout to dev channel and add

import os
os.environ["RWKV_JIT_ON"] = "1"

to your code before you init your model.

demo.zip