Is there a bug?

[zhirui-gao]

https://github.com/lifuguan/IGGT_official/blob/bec6604eb1972d1e8da8cf8315f773d9a321a25b/iggt/heads/window_sa.py#L275C1-L286C95

When executing the code, I encountered an error indicating that the tensor shapes are incorrect. Could you please explain the issue and suggest a fix?

[zhirui-gao]

def calculate_mask(self, x_size):
        h, w = x_size
        img_mask = torch.zeros((1, h, w, 1))
        h_slices = (slice(0, -self.window_size), slice(-self.window_size, None))
        w_slices = (slice(0, -self.window_size), slice(-self.window_size, None))
        cnt = 0
        for h_s in h_slices:
            for w_s in w_slices:
                img_mask[:, h_s, w_s, :] = cnt
                cnt += 1
        mask_windows = window_partition(img_mask, self.window_size)
        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
        return attn_mask

请问为什么我使用calculate_mask 得到的attn_mask全是0, 期待您的解答