remove code duplication

EleutherAI · Jul 28, 2023 · 02c2a9c · 02c2a9c
1 parent c397fc2
commit 02c2a9c
Showing 1 changed file with 10 additions and 19 deletions.
diff --git a/elk/training/ccs_reporter.py b/elk/training/ccs_reporter.py
@@ -100,25 +100,16 @@ def __init__(
 
  self.norm = None
 
- if self.config.norm == "burns":
- self.probe = nn.Sequential(
- nn.Linear(
- in_features,
- 1 if cfg.num_layers < 2 else hidden_size,
- bias=cfg.bias,
- device=device,
- ),
- nn.Sigmoid(),
- )
- else:
- self.probe = nn.Sequential(
- nn.Linear(
- in_features,
- 1 if cfg.num_layers < 2 else hidden_size,
- bias=cfg.bias,
- device=device,
- )
- )
+ self.probe = nn.Sequential(
+ nn.Linear(
+ in_features,
+ 1 if cfg.num_layers < 2 else hidden_size,
+ bias=cfg.bias,
+ device=device,
+ ),
+ *(nn.Sigmoid() if self.config.norm == "burns" else ()),
+ )
+
  if cfg.pre_ln:
  self.probe.insert(0, nn.LayerNorm(in_features, elementwise_affine=False))