Fix device mismatch and tensor.item() bugs

inference.py

@@ -129,13 +129,15 @@ class MultiHeadPredictor(nn.Module):
         Returns:
             Aggregated features [batch, seq, d_fiber]
         """
+        device = hidden_states[0].device
         fibers = []
         for i, (proj, hidden) in enumerate(zip(self.fiber_projs, hidden_states)):
             if i < len(hidden_states):
+                proj = proj.to(device)
                 fibers.append(proj(hidden.float()))
         
         # Weighted sum across layers
-        weights = F.softmax(self.layer_weights[:len(fibers)], dim=0)
+        weights = F.softmax(self.layer_weights.to(device)[:len(fibers)], dim=0)
         aggregated = sum(w * f for w, f in zip(weights, fibers))
         return aggregated
     
@@ -153,9 +155,11 @@ class MultiHeadPredictor(nn.Module):
         if not self.loaded_heads:
             return {}
         
+        device = hidden_states[0].device
         features = self.get_fiber_features(hidden_states)
         risks = {}
         for name in self.loaded_heads:
+            self.heads[name] = self.heads[name].to(device)
             logits = self.heads[name](features).squeeze(-1)
             risks[name] = torch.sigmoid(logits)
         return risks
@@ -370,28 +374,28 @@ class ARCSystem:
         interventions = {}
         
         # Repetition: suppress recently used tokens
-        if risks.get('repetition', torch.tensor(0)).item() > self.config.repetition_threshold:
+        if risks.get('repetition', 0) > self.config.repetition_threshold:
             for tok in set(recent_tokens[-self.config.repetition_window:]):
                 logits[0, tok] -= self.config.repetition_penalty
             interventions['repetition'] = True
             self.total_interventions['repetition'] += 1
         
         # Hedging: suppress hedge phrase starters
-        if risks.get('hedging', torch.tensor(0)).item() > self.config.hedging_threshold:
+        if risks.get('hedging', 0) > self.config.hedging_threshold:
             for tok in self._hedge_token_ids:
                 logits[0, tok] -= self.config.hedging_penalty
             interventions['hedging'] = True
             self.total_interventions['hedging'] += 1
         
         # Verbosity: suppress filler phrase starters
-        if risks.get('verbosity', torch.tensor(0)).item() > self.config.verbosity_threshold:
+        if risks.get('verbosity', 0) > self.config.verbosity_threshold:
             for tok in self._verbose_token_ids:
                 logits[0, tok] -= self.config.verbosity_penalty
             interventions['verbosity'] = True
             self.total_interventions['verbosity'] += 1
         
         # Sycophancy: suppress sycophantic starters
-        if risks.get('sycophancy', torch.tensor(0)).item() > self.config.sycophancy_threshold:
+        if risks.get('sycophancy', 0) > self.config.sycophancy_threshold:
             for tok in self._sycophancy_token_ids:
                 logits[0, tok] -= self.config.sycophancy_penalty
             interventions['sycophancy'] = True