Merge remote-tracking branch 'origin/main' into light-reporter-saving

elk/training/ccs_reporter.py

@@ -237,6 +237,17 @@ def reset_parameters(self):
 
  def forward(self, x: Tensor) -> Tensor:
  """Return the raw score output of the probe on `x`."""
+ assert x.shape[-2] == 2, "Probe input must be a contrast pair"
+
+ # Apply normalization
+ x0, x1 = x.unbind(-2)
+ x0, x1 = self.neg_norm(x0), self.pos_norm(x1)
+ x = torch.stack([x0, x1], dim=-2)
+
+ return self.raw_forward(x)
+
+ def raw_forward(self, x: Tensor) -> Tensor:
+ """Apply the probe to the provided input, without normalization."""
  return self.probe(x).mul(self.scale).add(self.bias).squeeze(-1)
 
  def loss(