Handling StopIteration error while generate activations

sparse_autoencoder/train/pipeline.py

@@ -1,35 +1,37 @@
 """Default pipeline."""
 from collections.abc import Iterator
 from functools import partial
+
+import itertools
 import logging
 from pathlib import Path
 from tempfile import gettempdir
 from typing import TYPE_CHECKING, final
 
 from jaxtyping import Float, Int
 from lightning import Trainer
 from lightning.pytorch.loggers import WandbLogger
 from pydantic import NonNegativeInt, PositiveInt, validate_call
 import torch
 from torch import Tensor
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformer_lens import HookedTransformer
 import wandb
 
 from sparse_autoencoder.activation_store.tensor_store import TensorActivationStore
 from sparse_autoencoder.autoencoder.lightning import LitSparseAutoencoder
 from sparse_autoencoder.metrics.validate.reconstruction_score import ReconstructionScoreMetric
 from sparse_autoencoder.metrics.wrappers.classwise import ClasswiseWrapperWithMean
 from sparse_autoencoder.source_data.abstract_dataset import SourceDataset, TorchTokenizedPrompts
 from sparse_autoencoder.source_model.replace_activations_hook import replace_activations_hook
 from sparse_autoencoder.source_model.store_activations_hook import store_activations_hook
 from sparse_autoencoder.source_model.zero_ablate_hook import zero_ablate_hook
 from sparse_autoencoder.train.utils.get_model_device import get_model_device
 from sparse_autoencoder.utils.data_parallel import DataParallelWithModelAttributes
 
 
 if TYPE_CHECKING:
     from sparse_autoencoder.tensor_types import Axis
 
 
@@ -139,7 +141,7 @@
         source_dataloader = source_dataset.get_dataloader(
             source_data_batch_size, num_workers=num_workers_data_loading
         )
-        self.source_data = iter(source_dataloader)
+        self.source_data = itertools.cycle(source_dataloader)
 
     @validate_call
     def generate_activations(self, store_size: PositiveInt) -> TensorActivationStore: