add @torch.no_grad to extract_hiddens

EleutherAI · Mar 15, 2023 · d50876e · d50876e
1 parent 2c6a10a
commit d50876e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/elk/extraction/extraction.py b/elk/extraction/extraction.py
@@ -71,6 +71,7 @@ def __post_init__(self, layer_stride: int):
             self.layers = tuple(range(0, config.num_hidden_layers, layer_stride))
 
 
+@torch.no_grad()
 def extract_hiddens(
     cfg: ExtractionConfig,
     *,