dev(narugo): try fix bug

.gitignore

@@ -251,3 +251,6 @@ $RECYCLE.BIN/
 
 # but keep examples
 !*.example
+/venv
+.python-version
+/.idea

tagger/model.py

@@ -125,9 +125,10 @@ def render_heatmap(
     },
     partial_rows: bool = True,
 ) -> tuple[list[Heatmap], Image.Image]:
-    hmap_dim = int(math.sqrt(pos_embed_dim))
+    # hmap_dim = int(math.sqrt(pos_embed_dim))
 
     image_hmaps = gradients.mean(2, keepdim=True).mul(image_feats.unsqueeze(0)).squeeze()
+    hmap_dim = int(math.sqrt(image_hmaps.mean(-1).shape[0] / len(image_labels)))
     image_hmaps = image_hmaps.mean(-1).reshape(len(image_labels), hmap_dim, hmap_dim)
     image_hmaps = image_hmaps.max(torch.zeros_like(image_hmaps))