ImportError

[KRL666]

Hello. I want to run rwkv_v6_demo.py but I'm encountering an error. How can I resolve it? Thank you.
Building extension module wkv6...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
[1/2] /usr/local/cuda-12.4/bin/nvcc --generate-dependencies-with-compile --dependency-output wkv6_cuda.cuda.o.d -DTORCH_EXTENSION_NAME=wkv6 -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE="gcc" -DPYBIND11_STDLIB="libstdcpp" -DPYBIND11_BUILD_ABI="cxxabi1011" -isystem /media/manbo/HHD/Anaconda3/envs/vrwkv/lib/python3.10/site-packages/torch/include -isystem /media/manbo/HHD/Anaconda3/envs/vrwkv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /media/manbo/HHD/Anaconda3/envs/vrwkv/lib/python3.10/site-packages/torch/include/TH -isystem /media/manbo/HHD/Anaconda3/envs/vrwkv/lib/python3.10/site-packages/torch/include/THC -isystem /usr/local/cuda-12.4/include -isystem /media/manbo/HHD/Anaconda3/envs/vrwkv/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS -D__CUDA_NO_HALF_CONVERSIONS_ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_89,code=compute_89 -gencode=arch=compute_89,code=sm_89 --compiler-options '-fPIC' -res-usage --use_fast_math -O3 --extra-device-vectorization -D_N_=64 -D_T_=4096 -std=c++17 -c /home/manbo/LKR/RWKV-LM-main/RWKV-v5/cuda/wkv6_cuda.cu -o wkv6_cuda.cuda.o
ptxas info : 3 bytes gmem, 24 bytes cmem[4]
ptxas info : Compiling entry function 'Z19kernel_backward_201IN3c108BFloat16EEviiiiPKT_S4_S4_S4_S4_S4_PS2' for 'sm_89'
ptxas info : Function properties for Z19kernel_backward_201IN3c108BFloat16EEviiiiPKT_S4_S4_S4_S4_S4_PS2
16376 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Used 96 registers, 16376 bytes cumulative stack size, 512 bytes smem, 424 bytes cmem[0]
ptxas info : Compiling entry function 'Z19kernel_backward_103IN3c108BFloat16EEviiiiPKT_S4_S4_S4_S4_S4_PS2' for 'sm_89'
ptxas info : Function properties for Z19kernel_backward_103IN3c108BFloat16EEviiiiPKT_S4_S4_S4_S4_S4_PS2
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Used 100 registers, 1024 bytes smem, 424 bytes cmem[0]
ptxas info : Compiling entry function 'Z19kernel_backward_102IN3c108BFloat16EEviiiiPKT_S4_S4_S4_S4_S4_PS2' for 'sm_89'
ptxas info : Function properties for Z19kernel_backward_102IN3c108BFloat16EEviiiiPKT_S4_S4_S4_S4_S4_PS2
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Used 98 registers, 512 bytes smem, 424 bytes cmem[0]
ptxas info : Compiling entry function 'Z19kernel_backward_101IN3c108BFloat16EEviiiiPKT_S4_S4_S4_S4_S4_PS2_S5' for 'sm_89'
ptxas info : Function properties for Z19kernel_backward_101IN3c108BFloat16EEviiiiPKT_S4_S4_S4_S4_S4_PS2_S5
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Used 119 registers, 512 bytes smem, 432 bytes cmem[0]
ptxas info : Compiling entry function 'Z14kernel_forwardIN3c108BFloat16EEviiiiPKT_S4_S4_S4_S4_PS2' for 'sm_89'
ptxas info : Function properties for Z14kernel_forwardIN3c108BFloat16EEviiiiPKT_S4_S4_S4_S4_PS2
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Used 100 registers, 1024 bytes smem, 416 bytes cmem[0]
[2/2] c++ wkv6_op.o wkv6_cuda.cuda.o -shared -L/media/manbo/HHD/Anaconda3/envs/vrwkv/lib/python3.10/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/usr/local/cuda-12.4/lib64 -lcudart -o wkv6.so
Loading extension module wkv6...
Traceback (most recent call last):
File "/home/manbo/LKR/RWKV-LM-main/RWKV-v5/src/rwkvp.py", line 271, in
wkv6_cuda = load(name="wkv6", sources=["RWKV-v5/cuda/wkv6_op.cpp", f"RWKV-v5/cuda/wkv6_cuda.cu"],
File "/media/manbo/HHD/Anaconda3/envs/vrwkv/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 1314, in load
return _jit_compile(
File "/media/manbo/HHD/Anaconda3/envs/vrwkv/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 1746, in _jit_compile
return _import_module_from_library(name, build_directory, is_python_module)
File "/media/manbo/HHD/Anaconda3/envs/vrwkv/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 2140, in _import_module_from_library
module = importlib.util.module_from_spec(spec)
File "", line 571, in module_from_spec
File "", line 1176, in create_module
File "", line 241, in _call_with_frames_removed
ImportError: dynamic module does not define module export function (PyInit_wkv6)

[BlinkDL]

pls update to latest torch 2.x and try again