From b14ea5642243abffa9da3dfe47838678b3bda763 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Wed, 22 Apr 2026 11:18:58 +0000
Subject: [PATCH] Backfill missing BatchNorm decoder state

Agent-Logs-Url: https://github.com/IBB666/IBB_POSE/sessions/39713ab6-656e-4be4-a016-26d6858e826a

Co-authored-by: kukulu00 <1170651+kukulu00@users.noreply.github.com>
---
 models/HeatmapHead.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/models/HeatmapHead.py b/models/HeatmapHead.py
index 35ece2c..507163e 100644
--- a/models/HeatmapHead.py
+++ b/models/HeatmapHead.py
@@ -107,7 +107,7 @@ class _FallbackUDPHeatmap:
 
 
 class _FallbackHeatmapHead(nn.Module):
-    _version = 2
+    _version = 3
 
     def __init__(
         self,
@@ -287,6 +287,19 @@ class _FallbackHeatmapHead(nn.Module):
 
             state_dict[prefix + new_key] = value
 
+        self._backfill_missing_batchnorm_state(state_dict, prefix)
+
+    def _backfill_missing_batchnorm_state(self, state_dict, prefix):
+        for module_name, module in self.named_modules():
+            if not isinstance(module, nn.BatchNorm2d):
+                continue
+
+            module_prefix = prefix + module_name + "."
+            for attr_name, attr_value in module.state_dict().items():
+                state_key = module_prefix + attr_name
+                if state_key not in state_dict:
+                    state_dict[state_key] = attr_value.detach().clone()
+
 
 def get_heatmap_head(mode="body"):
     head_cfg = _build_head_cfg(mode=mode)