RWKV_JIT_ON=0` is ignored when `RWKV_V7_ON=1

[Tornedahl]

Bug

When RWKV_V7_ON=1 is set, the RWKV_JIT_ON environment variable has no effect.
The V7 code path unconditionally sets MyModule = torch.jit.ScriptModule regardless
of the value of RWKV_JIT_ON.

Steps to reproduce

import os
os.environ["RWKV_V7_ON"] = "1"
os.environ["RWKV_JIT_ON"] = "0"
from rwkv.model import RWKV_x070 as RWKV
m = RWKV(model="your-model", strategy="cuda fp16")
print(hasattr(m, "w"))  # Returns False — expected True

Root cause

In model.py, lines 17–23 correctly check RWKV_JIT_ON:

if os.environ.get('RWKV_JIT_ON') != '0':
    MyModule = torch.jit.ScriptModule
else:
    MyModule = torch.nn.Module

But the RWKV_V7_ON block at line ~185 hardcodes:

MyModule = torch.jit.ScriptModule

...without checking the flag, overriding the earlier correct assignment.

Effect

model.w (the weight dict) becomes inaccessible when the model is compiled
as a ScriptModule. This makes in-process weight access impossible even when
the user has explicitly opted out of JIT.

Suggested fix

MyModule = torch.jit.ScriptModule if os.environ.get("RWKV_JIT_ON") != "0" else torch.nn.Module

Environment

• rwkv pip package (latest)
• RWKV-7 G1 model
• Python 3.12 / PyTorch 2.x

(Note: I discovered this while building an application on top of RWKV. I'm not a core contributor to the codebase, so please do verify — but the behavior is reproducible as described above.)

[BlinkDL]

Thank you :) Yeah the pip pkg is currently a mess. Please fork it and build your version if possible (put in https://github.com/BlinkDL/Albatross for significant speedup)