fix

lihuibng · Sep 16, 2023 · f5977f7 · f5977f7
1 parent e7a6e5a
commit f5977f7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/RWKV-v4neo/train.py b/RWKV-v4neo/train.py
@@ -162,7 +162,7 @@
     args.real_bsz = int(args.num_nodes) * int(args.devices) * args.micro_bsz
     os.environ["RWKV_T_MAX"] = str(args.ctx_len)
     os.environ["RWKV_MY_TESTING"] = args.my_testing
-    os.environ["RWKV_HEAD_SIZE_A"] = args.head_size_a
+    os.environ["RWKV_HEAD_SIZE_A"] = str(args.head_size_a)
     if args.dim_att <= 0:
         args.dim_att = args.n_embd
     if args.dim_ffn <= 0: