fix

pengzhouzp · Sep 16, 2023 · 8a6b4d2 · 8a6b4d2
1 parent f5977f7
commit 8a6b4d2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/RWKV-v4neo/src/model.py b/RWKV-v4neo/src/model.py
@@ -313,7 +313,7 @@ def forward(self, x):
 ########################################################################################################
 
 if 'r4' in os.environ["RWKV_MY_TESTING"]:
-    HEAD_SIZE = str(os.environ["RWKV_HEAD_SIZE_A"])
+    HEAD_SIZE = int(os.environ["RWKV_HEAD_SIZE_A"])
     wkv5_cuda = load(name="wkv5", sources=["cuda/wkv5_op.cpp", f"cuda/wkv5_cuda.cu"],
                     verbose=True, extra_cuda_cflags=["-res-usage", "--use_fast_math", "-O3", "-Xptxas -O3", "--extra-device-vectorization", f"-D_N_={HEAD_SIZE}"])