replace slicing

EleutherAI · Aug 11, 2023 · af4c9d1 · af4c9d1
1 parent cc095cb
commit af4c9d1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/elk/extraction/extraction.py b/elk/extraction/extraction.py
@@ -194,7 +194,7 @@ def extract_hiddens(
  seed=cfg.seed,
  )
 
- layer_indices = cfg.layers or tuple(range(model.config.num_hidden_layers)[1:])
+ layer_indices = cfg.layers or tuple(range(1, model.config.num_hidden_layers))
 
  global_max_examples = cfg.max_examples[0 if split_type == "train" else 1]
 
@@ -366,7 +366,7 @@ def hidden_features(cfg: Extract) -> tuple[DatasetInfo, Features]:
  if num_dropped:
  print(f"Dropping {num_dropped} non-multiple choice templates")
 
- layer_indices = cfg.layers or tuple(range(model_cfg.num_hidden_layers)[1:])
+ layer_indices = cfg.layers or tuple(range(1, model_cfg.num_hidden_layers))
  layer_cols = {
  f"hidden_{layer}": Array3D(
  dtype="int16",