opt ppo model

src/model.py

@@ -20,6 +20,7 @@ from src.rlhf.utils import exists
 from src.rlhf.utils import gumbel_sample
 from src.rlhf.utils import top_k
 from src.rlhf.utils import identity
+from src.rlhf.utils import eval_decorator
 
 # from deepspeed.runtime.fp16.onebit.zoadam import ZeroOneAdam
 
@@ -493,6 +494,7 @@ class RWKV(pl.LightningModule):
         return logits
     
     @torch.no_grad()
+    @eval_decorator
     def generate(
         self,
         seq_len,