Update core/profiler.py

core/profiler.py

@@ -7,23 +7,35 @@ class ActivationProfiler:
         self.stats = defaultdict(dict)
         self.hooks = []
 
-    def hook_fn(self, name):
-        def hook(module, input, output):
-            act = output.detach().float()
-            std = act.std().item()
-            mean = act.mean().item()
-            min_val = act.min().item()
-            max_val = act.max().item()
-            outliers = (act.abs() > 3 * std).float().mean().item()
-
-            self.stats[name] = {
-                "mean": mean,
-                "std": std,
-                "min": min_val,
-                "max": max_val,
-                "outlier_ratio": outliers
-            }
-        return hook
+
+    def hook(module, input, output):
+
+        # 🔥 FIX: Handle tuple outputs safely
+        if isinstance(output, tuple):
+            output = output[0]
+
+        if not isinstance(output, torch.Tensor):
+            return
+
+        act = output.detach().float()
+
+        std = act.std().item()
+        mean = act.mean().item()
+        min_val = act.min().item()
+        max_val = act.max().item()
+
+        outliers = (act.abs() > 3 * std).float().mean().item()
+
+        self.stats[name] = {
+            "mean": mean,
+            "std": std,
+            "min": min_val,
+            "max": max_val,
+            "outlier_ratio": outliers
+        }
+
+    return hook
+
 
     def register(self):
         for name, module in self.model.named_modules():