From 7476c69f324250f55ac6bfa9745540645b72c2d3 Mon Sep 17 00:00:00 2001
From: BlinkDL <a@a.com>
Date: Sun, 12 Feb 2023 09:28:26 +0000
Subject: [PATCH] fix

---
 RWKV-v4neo/src/model.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/RWKV-v4neo/src/model.py b/RWKV-v4neo/src/model.py
index 7635b6c..621a398 100644
--- a/RWKV-v4neo/src/model.py
+++ b/RWKV-v4neo/src/model.py
@@ -14,6 +14,10 @@ from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
 
 # from deepspeed.runtime.fp16.onebit.zoadam import ZeroOneAdam
 
+try:
+    print('RWKV_MY_TESTING', os.environ["RWKV_MY_TESTING"])
+except:
+    os.environ["RWKV_MY_TESTING"] = ''
 
 def __nop(ob):
     return ob
@@ -346,6 +350,14 @@ class RWKV(pl.LightningModule):
     def __init__(self, args):
         super().__init__()
         self.args = args
+        if not hasattr(args, 'dim_att'):
+            args.dim_att = args.n_embd
+        if not hasattr(args, 'dim_ffn'):
+            args.dim_ffn = args.n_embd * 4
+        if not hasattr(args, 'tiny_att_layer'):
+            args.tiny_att_layer = -1
+        if not hasattr(args, 'tiny_att_dim'):
+            args.tiny_att_dim = -1
 
         self.emb = nn.Embedding(args.vocab_size, args.n_embd)