simplified padding

Signed-off-by: Dan <Dan@DESKTOP-UQ7GDK9>
RWKV · Jan 31, 2024 · fd0342f · fd0342f
1 parent ded473e
commit fd0342f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/RWKV-v5/src/module/rwkv_inner.py b/RWKV-v5/src/module/rwkv_inner.py
@@ -49,7 +49,7 @@ def rwkv_inner(r,k,v,w,u,kv_state,chunk_len:int=24,precision_dtype:torch.dtype=t
         wc_log_cum = wc_log.cumsum(dim=-2)
 
         # chunked view of shifted_w_log
-        shifted_wc_log_cum = torch.cat([torch.zeros_like(wc_log_cum[:,:,:,:1,:]), wc_log_cum[:,:,:,:-1,:]], dim=-2)
+        shifted_wc_log_cum = F.pad(wc_log_cum, (0, 0, 1, -1))
 
 
         # NOTE - we have to apply the decay weight from TWO ahead.. ONE ahead gets no decay (log==0)