Update core/profiler.py

core/profiler.py

@@ -8,33 +8,35 @@ class ActivationProfiler:
         self.hooks = []
 
 
-    def hook(module, input, output):
+    def hook_fn(self, name):
+        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
 
-        # 🔥 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):