fix(training): stable GPU usage

OpenMOSS · Jul 30, 2024 · 0d333e1 · 0d333e1
1 parent a4e1c90
commit 0d333e1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/lm_saes/runner.py b/src/lm_saes/runner.py
@@ -328,7 +328,7 @@ def sample_feature_activations_runner(cfg: LanguageModelSAEAnalysisConfig):
         sae = parallelize_module(sae, device_mesh=sae.device_mesh["tp"], parallelize_plan=plan) # type: ignore
         sae.parallelize_plan = plan
 
-    sae.decoder.weight = None
+    sae.decoder.weight = None  # type: ignore[assignment]
     torch.cuda.empty_cache()