fix ruff check

Author: Hzfinfdu

src/lm_saes/evaluator.py

@@ -4,10 +4,10 @@
 from torch import Tensor
 from tqdm import tqdm
 from transformer_lens import HookedTransformer
+from wandb.sdk.wandb_run import Run
 
 from lm_saes.config import EvalConfig
 from lm_saes.sae import SparseAutoEncoder
-from wandb.sdk.wandb_run import Run
 
 
 class Evaluator:

src/lm_saes/runner.py

@@ -2,11 +2,11 @@
 from pathlib import Path
 from typing import Literal, Optional, TypeVar, overload
 
+import wandb
 from pydantic import model_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from torch.distributed.device_mesh import init_device_mesh
 
-import wandb
 from lm_saes.activation.factory import ActivationFactory
 from lm_saes.activation.writer import ActivationWriter
 from lm_saes.analysis.feature_analyzer import FeatureAnalyzer

src/lm_saes/sae.py

@@ -109,6 +109,7 @@ def topk_activation(x: torch.Tensor):
                 k = x.shape[-1] - self.current_k + 1
                 k_th_value, _ = torch.kthvalue(x, k=k, dim=-1)
                 k_th_value = k_th_value.unsqueeze(dim=1)
+                print()
                 return x.ge(k_th_value)
 
             return topk_activation

src/lm_saes/trainer.py

@@ -7,12 +7,12 @@
 from torch import Tensor
 from torch.optim import Adam, Optimizer
 from tqdm import tqdm
+from wandb.sdk.wandb_run import Run
 
 from lm_saes.config import TrainerConfig
 from lm_saes.optim import get_scheduler
 from lm_saes.sae import SparseAutoEncoder
 from lm_saes.utils.misc import all_reduce_tensor
-from wandb.sdk.wandb_run import Run
 
 
 class Trainer: