Fixing issue in fast_pos_embed

Signed-off-by: vtirumal <vtirumal@qti.qualcomm.com>

Author: tv-karthikeya

QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py

@@ -192,23 +192,20 @@ def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
         return embeddings
 
     def fast_pos_embed_interpolate(self, grid_thw):
-        # breakpoint()
-        # gridbs, grid_ts, grid_hs, grid_ws = grid_thw.shape
-        bs, t, h, w = grid_thw.shape
-        # grid_ts = torch.tensor([grid_ts], device=grid_thw.device)
-        # grid_hs = torch.tensor([grid_hs], device=grid_thw.device)
-        # grid_ws = torch.tensor([grid_ws], device=grid_thw.device)
-        idx_list = [[] for _ in range(4)]
-        weight_list = [[] for _ in range(4)]
-        # t,h,w = grid_ts[0],grid_hs[0],grid_ws[0]
-        # for t, h, w in zip(grid_ts, grid_hs, grid_ws):
+        bs,t,h,w=grid_thw.shape
+
         h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h)
         w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w)
 
         h_idxs_floor = h_idxs.int()
         w_idxs_floor = w_idxs.int()
-        h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
-        w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
+        # h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
+        # w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
+        # TO resolve clip issue
+        max_t = torch.tensor(self.num_grid_per_side - 1, device=h_idxs.device)
+
+        h_idxs_ceil = torch.minimum(h_idxs_floor + 1, max_t) # working
+        w_idxs_ceil = torch.minimum(w_idxs_floor + 1, max_t)
 
         dh = h_idxs - h_idxs_floor
         dw = w_idxs - w_idxs_floor
@@ -230,12 +227,12 @@ def fast_pos_embed_interpolate(self, grid_thw):
             (dh[None].T * dw[None]).flatten(),
         ]
 
-        for i in range(4):
-            idx_list[i].extend(indices[i].tolist())
-            weight_list[i].extend(weights[i].tolist())
-        idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=self.pos_embed.weight.device)
-        weight_tensor = torch.tensor(
-            weight_list, dtype=self.pos_embed.weight.dtype, device=self.pos_embed.weight.device
+        idx_tensor = torch.stack(indices, dim=0).to(
+            dtype=torch.long, device=self.pos_embed.weight.device
+        )  # [4, h*w]
+
+        weight_tensor = torch.stack(weights, dim=0).to(
+            dtype=self.pos_embed.weight.dtype, device=self.pos_embed.weight.device
         )
         pos_embeds = self.pos_embed(idx_tensor) * weight_tensor[:, :, None]
         patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]
@@ -244,9 +241,7 @@ def fast_pos_embed_interpolate(self, grid_thw):
 
         patch_pos_embeds_permute = []
         merge_size = self.config.spatial_merge_size
-        # breakpoint()
         pos_embed = patch_pos_embeds[0]
-        # for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
         pos_embed = pos_embed.repeat(t, 1)
 
         pos_embed = (