fix assignment of probe

EleutherAI · Aug 3, 2023 · f2ae1b7 · f2ae1b7
1 parent aeb8b09
commit f2ae1b7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/elk/training/ccs_reporter.py b/elk/training/ccs_reporter.py
@@ -99,7 +99,7 @@ def __init__(
  hidden_size = cfg.hidden_size or 4 * in_features // 3
 
  self.norm = None
- self.probe: nn.Sequential(
+ self.probe = nn.Sequential(
  nn.Linear(
  in_features,
  1 if cfg.num_layers < 2 else hidden_size,