Update moondream.py

moondream.py

@@ -609,8 +609,8 @@ class MoondreamModel(nn.Module):
             )
             inputs_embeds = torch.cat([bos_emb, img_emb[None]], dim=1)
             attn = self.attn_mask          # (1,1,Tmax,Tmax)
-            mask = attn[:, :, pos:pos+T, :].expand(B, -1, -1, -1).contiguous()  # (B,1,T,K)
-            pos_ids = torch.arange(pos, pos+T, device=self.device, dtype=torch.long)
+            mask = self.attn_mask[:, :, 0 : inputs_embeds.size(1), :]
+            pos_ids = torch.arange(inputs_embeds.size(1), dtype=torch.long)
             self._prefill(inputs_embeds, mask, pos_ids, lora)