fix pyright

elk/run.py

@@ -102,7 +102,7 @@ def execute(
 
  devices = select_usable_devices(self.num_gpus, min_memory=self.min_gpu_mem)
  num_devices = len(devices)
- func: Callable[[int], dict[str, pd.DataFrame]] = partial(
+ func: Callable[[int], list[dict[str, pd.DataFrame]]] = partial(
  self.apply_to_layer,
  devices=devices,
  world_size=num_devices,
@@ -113,7 +113,7 @@ def execute(
  @abstractmethod
  def apply_to_layer(
  self, layer: int, devices: list[str], world_size: int, probe_per_prompt: bool
- ) -> dict[str, pd.DataFrame]:
+ ) -> list[dict[str, pd.DataFrame]]:
  """Train or eval a reporter on a single layer."""
 
  def make_reproducible(self, seed: int):
@@ -162,7 +162,7 @@ def concatenate(self, layers):
 
  def apply_to_layers(
  self,
- func: Callable[[int], dict[str, pd.DataFrame]],
+ func: Callable[[int], list[dict[str, pd.DataFrame]]],
  num_devices: int,
  ):
  """Apply a function to each layer of the datasets in parallel

elk/training/train.py

@@ -232,6 +232,8 @@ def apply_to_layer(
  probe_per_prompt: bool,
  ) -> list[dict[str, pd.DataFrame]]:
  """Train a single reporter on a single layer."""
+ assert self.out_dir is not None # TODO this is really annoying, why can it be
+ # None?
 
  self.make_reproducible(seed=self.net.seed + layer)
  device = self.get_device(devices, world_size)