From fd2c196b48be2124e0663bf48e2acafb939eb36f Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 17 Jan 2025 17:10:56 +0200
Subject: [PATCH] Update posemb_layers.py

---
 hyvideo/modules/posemb_layers.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/hyvideo/modules/posemb_layers.py b/hyvideo/modules/posemb_layers.py
index fa42dbe..2ac4ce7 100644
--- a/hyvideo/modules/posemb_layers.py
+++ b/hyvideo/modules/posemb_layers.py
@@ -92,7 +92,8 @@ def apply_rotary_emb(
 
     """
    
-    cos, sin = [f.view(*xq.shape[:2], 1, xq.shape[3]) for f in freqs_cis]
+    shape = [d if i == 1 or i == xq.ndim - 1 else 1 for i, d in enumerate(xq.shape)]
+    cos, sin = freqs_cis[0].view(*shape), freqs_cis[1].view(*shape)
 
     if upcast:
         xq_out = apply_rotary(xq.float(), cos, sin).to(xq.dtype)