From c819ae0ffac76290f367254299bc6115822b54d3 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Tue, 29 Jul 2025 10:25:32 +0200
Subject: [PATCH 01/21] added option to specify list of module names instead of
 single name

---
 thingsvision/core/extraction/base.py       | 174 ++++++++++++---------
 thingsvision/core/extraction/tensorflow.py |  30 +++-
 thingsvision/core/extraction/torch.py      | 130 ++++++++-------
 3 files changed, 200 insertions(+), 134 deletions(-)

diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index e1bf3fe5..4c67409b 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -2,7 +2,7 @@
 import os
 import re
 import warnings
-from typing import Callable, Iterator, List, Optional, Union
+from typing import Callable, Dict, Iterator, List, Optional, Union
 
 import numpy as np
 from torchtyping import TensorType
@@ -76,17 +76,21 @@ def load_model(self) -> None:
     def extract_batch(
         self,
         batch: Union[TensorType["b", "c", "h", "w"], Array],
-        module_name: str,
+        module_name: Optional[str] = None,
+        module_names: Optional[List[str]] = None,
         flatten_acts: bool,
         output_type: str,
-    ) -> Union[
+    ) -> Dict[
+        str,
         Union[
-            TensorType["b", "num_maps", "h_prime", "w_prime"],
-            TensorType["b", "t", "d"],
-            TensorType["b", "p"],
-            TensorType["b", "d"],
+            Union[
+                TensorType["n", "num_maps", "h_prime", "w_prime"],
+                TensorType["n", "t", "d"],
+                TensorType["n", "p"],
+                TensorType["n", "d"],
+            ],
+            Array,
         ],
-        Array,
     ]:
         """Extract the activations of a selected module for every image in a mini-batch.
 
@@ -94,8 +98,8 @@ def extract_batch(
         ----------
         batch : np.ndarray or torch.Tensor
             mini-batch of three-dimensional image tensors.
-        module_name : str
-            Name of the module for which features should be extraced.
+        module_names : List[str]
+            Names of the modules for which features should be extracted.
         flatten_acts : bool
             Whether the activation of a tensor should be flattened to a vector.
         output_type : str {"ndarray", "tensor"}
@@ -112,16 +116,20 @@ def extract_batch(
     def _extract_batch(
         self,
         batch: Union[TensorType["b", "c", "h", "w"], Array],
-        module_name: str,
+        module_name: Optional[str] = None,
+        module_names: Optional[List[str]] = None,
         flatten_acts: bool,
-    ) -> Union[
+    ) -> Dict[
+        str,
         Union[
-            TensorType["b", "num_maps", "h_prime", "w_prime"],
-            TensorType["b", "t", "d"],
-            TensorType["b", "p"],
-            TensorType["b", "d"],
+            Union[
+                TensorType["n", "num_maps", "h_prime", "w_prime"],
+                TensorType["n", "t", "d"],
+                TensorType["n", "p"],
+                TensorType["n", "d"],
+            ],
+            Array,
         ],
-        Array,
     ]:
         raise NotImplementedError
 
@@ -129,13 +137,16 @@ def get_output_types(self) -> List[str]:
         """Return the list of available output types (for the feature matrix)."""
         return ["ndarray", "tensor"]
 
-    def _module_and_output_check(self, module_name: str, output_type: str) -> None:
+    def _module_and_output_check(
+        self, module_names: List[str], output_type: str
+    ) -> None:
         """Checks whether the provided module name and output type are valid."""
         valid_names = self.get_module_names()
-        if not module_name in valid_names:
-            raise ValueError(
-                f"\n{module_name} is not a valid module name. Please choose a name from the following set of modules: {valid_names}\n"
-            )
+        for module_name in module_names:
+            if module_name not in valid_names:
+                raise ValueError(
+                    f"\n{module_name} is not a valid module name. Please choose a name from the following set of modules: {valid_names}\n"
+                )
         assert (
             output_type in self.get_output_types()
         ), f"\nData type of output feature matrix must be set to one of the following available data types: {self.get_output_types()}\n"
@@ -143,19 +154,23 @@ def _module_and_output_check(self, module_name: str, output_type: str) -> None:
     def extract_features(
         self,
         batches: Iterator[Union[TensorType["b", "c", "h", "w"], Array]],
-        module_name: str,
+        module_name: Optional[str] = None,
+        module_names: Optional[List[str]] = None,
         flatten_acts: bool = False,
         output_type: Optional[str] = "ndarray",
         output_dir: Optional[str] = None,
         step_size: Optional[int] = None,
-    ) -> Union[
+    ) -> Dict[
+        str,
         Union[
-            TensorType["n", "num_maps", "h_prime", "w_prime"],
-            TensorType["n", "t", "d"],
-            TensorType["n", "p"],
-            TensorType["n", "d"],
+            Union[
+                TensorType["n", "num_maps", "h_prime", "w_prime"],
+                TensorType["n", "t", "d"],
+                TensorType["n", "p"],
+                TensorType["n", "d"],
+            ],
+            Array,
         ],
-        Array,
     ]:
         """Extract hidden unit activations (at specified layer) for every image in the database.
 
@@ -166,8 +181,11 @@ def extract_features(
             mini-batches, where each element is a
             subsample of the full (image) dataset.
         module_name : str
-            Layer name. Name of neural network layer for
-            which features should be extraced.
+            Layer name. Name of the neural network layer for
+            which features should be extracted.
+        module_names : List[str]
+            Layer names. Names of neural network layers for
+            which features should be extracted.
         flatten_acts : bool
             Whether activation tensor (e.g., activations
             from an early layer of the neural network model)
@@ -195,7 +213,12 @@ def extract_features(
         output : np.ndarray or torch.Tensor
             Returns the feature matrix (e.g., $X \in \mathbb{R}^{n \times d}$ if penultimate or logits layer or flatten_acts = True).
         """
-        self._module_and_output_check(module_name, output_type)
+        assert (
+            module_name ^ module_names
+        ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
+        if module_name:
+            module_names = [module_name]
+        self._module_and_output_check(module_names, output_type)
         if output_dir:
             os.makedirs(output_dir, exist_ok=True)
 
@@ -203,44 +226,47 @@ def extract_features(
                 # if step size is not given, assume that features to every image consume 3MB of memory and that the user has at least 8GB of free RAM
                 step_size = 8000 // (len(next(iter(batches))) * 3) + 1
 
-        features = []
+        # create feature dict per module name
+        features = {name: [] for name in module_names}
         image_ct, last_image_ct = 0, 0
         for i, batch in tqdm(
             enumerate(batches, start=1), desc="Batch", total=len(batches)
         ):
-            features.append(
-                self._extract_batch(
-                    batch=batch, module_name=module_name, flatten_acts=flatten_acts
-                )
+            modules_features = self._extract_batch(
+                batch=batch, module_names=module_names, flatten_acts=flatten_acts
             )
 
             image_ct += len(batch)
             del batch
 
-            if output_dir and (i % step_size == 0 or i == len(batches)):
-                if self.get_backend() == "pt":
-                    features_subset = torch.cat(features)
-                    if output_type == "ndarray":
-                        features_subset = self._to_numpy(features_subset)
+            for module_name in module_names:
+                features[module_name].append(modules_features[module_name])
+
+                if output_dir and (i % step_size == 0 or i == len(batches)):
+                    if self.get_backend() == "pt":
+                        features_subset = torch.cat(features[module_name])
+                        if output_type == "ndarray":
+                            features_subset = self._to_numpy(features_subset)
+                            features_subset_file = os.path.join(
+                                output_dir,
+                                f"features_{module_name}_{last_image_ct}-{image_ct}.npy",
+                            )
+                            np.save(features_subset_file, features_subset)
+                        else:  # output_type = tensor
+                            features_subset_file = os.path.join(
+                                output_dir,
+                                f"features_{module_name}_{last_image_ct}-{image_ct}.pt",
+                            )
+                            torch.save(features_subset, features_subset_file)
+                    else:
                         features_subset_file = os.path.join(
                             output_dir,
-                            f"features_{last_image_ct}-{image_ct}.npy",
+                            f"features_{module_name}_{last_image_ct}-{image_ct}.npy",
                         )
+                        features_subset = np.vstack(features[module_name])
                         np.save(features_subset_file, features_subset)
-                    else:  # output_type = tensor
-                        features_subset_file = os.path.join(
-                            output_dir,
-                            f"features_{last_image_ct}-{image_ct}.pt",
-                        )
-                        torch.save(features_subset, features_subset_file)
-                else:
-                    features_subset_file = os.path.join(
-                        output_dir, f"features_{last_image_ct}-{image_ct}.npy"
-                    )
-                    features_subset = np.vstack(features)
-                    np.save(features_subset_file, features_subset)
-                features = []
-                last_image_ct = image_ct
+                    features = {name: [] for name in module_names}
+                    last_image_ct = image_ct
         print(
             f"...Features successfully extracted for all {image_ct} images in the database."
         )
@@ -248,26 +274,30 @@ def extract_features(
             print(f"...Features were saved to {output_dir}.")
             return None
         else:
-            if self.get_backend() == "pt":
-                features = torch.cat(features)
-                if output_type == "ndarray":
-                    features = self._to_numpy(features)
-            else:
-                features = np.vstack(features)
-                print(f"...Features shape: {features.shape}")
+            for module_name in module_names:
+                if self.get_backend() == "pt":
+                    features = torch.cat(features[module_name])
+                    if output_type == "ndarray":
+                        features = self._to_numpy(features)
+                else:
+                    features = np.vstack(features[module_name])
+                    print(f"...Features shape: {features.shape}")
         return features
 
     @staticmethod
     def _to_numpy(
-        features: Union[
-            TensorType["n", "num_maps", "h_prime", "w_prime"],
-            TensorType["n", "t", "d"],
-            TensorType["n", "p"],
-            TensorType["n", "d"],
-        ]
-    ) -> Array:
+        features: Dict[
+            str,
+            Union[
+                TensorType["n", "num_maps", "h_prime", "w_prime"],
+                TensorType["n", "t", "d"],
+                TensorType["n", "p"],
+                TensorType["n", "d"],
+            ],
+        ],
+    ) -> Dict[str, Array]:
         """Move activations to CPU and convert torch.Tensor to np.ndarray."""
-        return features.numpy()
+        return {k: v.cpu().numpy() for k, v in features.items()}
 
     def get_transformations(
         self, resize_dim: int = 256, crop_dim: int = 224, apply_center_crop: bool = True
diff --git a/thingsvision/core/extraction/tensorflow.py b/thingsvision/core/extraction/tensorflow.py
index 1dc49f53..31957710 100644
--- a/thingsvision/core/extraction/tensorflow.py
+++ b/thingsvision/core/extraction/tensorflow.py
@@ -38,13 +38,19 @@ def __init__(
     def _extract_batch(
         self,
         batch: Array,
-        module_name: str,
+        module_name: Optional[str] = None,
+        module_names: Optional[List[str]] = None,
         flatten_acts: bool,
     ) -> Array:
-        layer_out = [self.model.get_layer(module_name).output]
+        assert (
+            module_name ^ module_names
+        ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
+        if module_name:
+            module_names = [module_name]
+        layer_outs = [self.model.get_layer(name).output for name in module_names]
         activation_model = keras.models.Model(
             inputs=self.model.input,
-            outputs=layer_out,
+            outputs=layer_outs,
         )
         activations = activation_model.predict(batch)
         if flatten_acts:
@@ -54,12 +60,22 @@ def _extract_batch(
     def extract_batch(
         self,
         batch: Array,
-        module_name: str,
-        flatten_acts: bool,
+        module_name: Optional[str] = None,
+        module_names: Optional[List[str]] = None,
+        flatten_acts: bool = False,
         output_type: str = "ndarray",
     ) -> Array:
-        self._module_and_output_check(module_name, output_type)
-        activations = self._extract_batch(batch, module_name, flatten_acts)
+        assert (
+            module_name ^ module_names
+        ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
+        if module_name:
+            module_names = [module_name]
+        self.model = self.model.to(self.device)
+        self.activations = {}
+        if module_name:
+            module_names = [module_name]
+        self._module_and_output_check(module_names, output_type)
+        activations = self._extract_batch(batch, module_names, flatten_acts)
         return activations
 
     def show_model(self) -> str:
diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
index 3f854ef0..3dcf34f2 100644
--- a/thingsvision/core/extraction/torch.py
+++ b/thingsvision/core/extraction/torch.py
@@ -73,21 +73,30 @@ def hook(model, input, output) -> None:
 
         return hook
 
-    def _register_hook(self, module_name: str) -> None:
-        """Register a forward hook to store activations."""
+    def _register_hooks(self, module_names: List[str]) -> None:
+        """Register a forward hook to store activations for multiple modules."""
+        self.hook_handles = []
         for n, m in self.model.named_modules():
-            if n == module_name:
-                self.hook_handle = m.register_forward_hook(self.get_activation(n))
-                break
-
-    def _unregister_hook(self) -> None:
-        """Remove the forward hook."""
-        self.hook_handle.remove()
+            if n in module_names:
+                handle = m.register_forward_hook(self.get_activation(n))
+                self.hook_handles.append(handle)
+
+    def _unregister_hooks(self) -> None:
+        """Unregister all forward hooks."""
+        if self.hook_handles:
+            for handle in self.hook_handles:
+                handle.remove()
+            self.hook_handles = []
+        else:
+            warnings.warn(
+                "\nNo hooks were registered. Nothing to unregister.\n",
+                category=UserWarning,
+            )
 
-    def batch_extraction(self, module_name: str, output_type: str) -> object:
+    def batch_extraction(self, module_names: List[str], output_type: str) -> object:
         """Allows mini-batch extraction for custom data pipeline using a with-statement."""
         return BatchExtraction(
-            extractor=self, module_name=module_name, output_type=output_type
+            extractor=self, module_names=module_names, output_type=output_type
         )
 
     def extract_batch(
@@ -101,7 +110,7 @@ def extract_batch(
         TensorType["b", "d"],
     ]:
         act = self._extract_batch(
-            batch=batch, module_name=self.module_name, flatten_acts=flatten_acts
+            batch=batch, module_name=self.module_names, flatten_acts=flatten_acts
         )
         if self.output_type == "ndarray":
             act = self._to_numpy(act)
@@ -111,48 +120,55 @@ def extract_batch(
     def _extract_batch(
         self,
         batch: TensorType["b", "c", "h", "w"],
-        module_name: str,
+        module_names: List[str],
         flatten_acts: bool,
-    ) -> Union[
-        TensorType["b", "num_maps", "h_prime", "w_prime"],
-        TensorType["b", "t", "d"],
-        TensorType["b", "p"],
-        TensorType["b", "d"],
+    ) -> Dict[
+        str,
+        Union[
+            TensorType["b", "num_maps", "h_prime", "w_prime"],
+            TensorType["b", "t", "d"],
+            TensorType["b", "p"],
+            TensorType["b", "d"],
+        ],
     ]:
         """Extract representations from a batch of images."""
         # move mini-batch to current device
         batch = batch.to(self.device)
         _ = self.forward(batch)
-        act = self.activations[module_name]
-        if len(act.shape) > 2:
-            if hasattr(self, "token_extraction"):
-                if self.token_extraction == "cls_token":
-                    act = act[:, 0, :].clone()
-                elif self.token_extraction == "avg_pool":
-                    act = act[:, 1:, :].clone().mean(dim=1)
-                elif self.token_extraction == "cls_token+avg_pool":
-                    cls_token = act[:, 0, :].clone()
-                    pooled_tokens = act[:, 1:, :].clone().mean(dim=1)
-                    act = torch.cat((cls_token, pooled_tokens), dim=1)
-                else:
-                    raise ValueError(
-                        f"\n{self.token_extraction} is not a valid value for token extraction. "
-                        "\nChoose one of the following: {TOKEN_EXTRACTIONS}.\n "
-                    )
-            elif flatten_acts:
-                if self.model_name.lower().startswith("clip"):
-                    act = self.flatten_acts(act, batch, module_name)
-                else:
-                    act = self.flatten_acts(act)
-        if act.is_cuda or act.get_device() >= 0:
-            torch.cuda.empty_cache()
-            act = act.cpu()
-        return act
+        acts = {}
+        for module_name in module_names:
+            act = self.activations[module_name]
+            if len(act.shape) > 2:
+                if hasattr(self, "token_extraction"):
+                    if self.token_extraction == "cls_token":
+                        act = act[:, 0, :].clone()
+                    elif self.token_extraction == "avg_pool":
+                        act = act[:, 1:, :].clone().mean(dim=1)
+                    elif self.token_extraction == "cls_token+avg_pool":
+                        cls_token = act[:, 0, :].clone()
+                        pooled_tokens = act[:, 1:, :].clone().mean(dim=1)
+                        act = torch.cat((cls_token, pooled_tokens), dim=1)
+                    else:
+                        raise ValueError(
+                            f"\n{self.token_extraction} is not a valid value for token extraction. "
+                            "\nChoose one of the following: {TOKEN_EXTRACTIONS}.\n "
+                        )
+                elif flatten_acts:
+                    if self.model_name.lower().startswith("clip"):
+                        act = self.flatten_acts(act, batch, module_name)
+                    else:
+                        act = self.flatten_acts(act)
+            if act.is_cuda or act.get_device() >= 0:
+                torch.cuda.empty_cache()
+                act = act.cpu()
+            acts[module_name] = act
+        return acts
 
     def extract_features(
         self,
         batches: Iterator,
-        module_name: str,
+        module_name: Optional[str] = None,
+        module_names: Optional[List[str]] = None,
         flatten_acts: bool = False,
         output_type: str = "ndarray",
         output_dir: Optional[str] = None,
@@ -160,17 +176,21 @@ def extract_features(
     ):
         self.model = self.model.to(self.device)
         self.activations = {}
-        self._register_hook(module_name=module_name)
+        if module_name:
+            module_names = [module_name]
+        self._register_hooks(module_names=module_names)
         features = super().extract_features(
             batches=batches,
             module_name=module_name,
+            module_names=module_names,
             flatten_acts=flatten_acts,
             output_type=output_type,
             output_dir=output_dir,
             step_size=step_size,
         )
-        if self.hook_handle:
-            self._unregister_hook()
+        if self.hook_handles:
+            for handle in self.hook_handles:
+                handle.remove()
         return features
 
     def forward(
@@ -183,7 +203,7 @@ def forward(
     def flatten_acts(
         act: Union[
             TensorType["b", "num_maps", "h_prime", "w_prime"], TensorType["b", "t", "d"]
-        ]
+        ],
     ) -> TensorType["b", "p"]:
         """Default flattening of activations."""
         return act.view(act.size(0), -1)
@@ -270,7 +290,7 @@ def get_backend(self) -> str:
 class BatchExtraction(object):
 
     def __init__(
-        self, extractor: PyTorchExtractor, module_name: str, output_type: str
+        self, extractor: PyTorchExtractor, module_names: List[str], output_type: str
     ) -> None:
         """
         Mini-batch extraction object that can be used as a with-statement in a PyTorch extractor.
@@ -283,19 +303,19 @@ def __init__(
 
         """
         self.extractor = extractor
-        self.module_name = module_name
+        self.module_names = module_names
         self.output_type = output_type
 
     def __enter__(self) -> PyTorchExtractor:
         """Registering hooks and setting attributes during opening."""
-        self.extractor._module_and_output_check(self.module_name, self.output_type)
-        self.extractor._register_hook(self.module_name)
-        setattr(self.extractor, "module_name", self.module_name)
+        self.extractor._module_and_output_check(self.module_names, self.output_type)
+        self.extractor._register_hooks(self.module_names)
+        setattr(self.extractor, "module_names", self.module_names)
         setattr(self.extractor, "output_type", self.output_type)
         return self.extractor
 
     def __exit__(self, *args):
         """Removing hooks and deleting attributes at closing."""
-        self.extractor._unregister_hook()
-        delattr(self.extractor, "module_name")
+        self.extractor._unregister_hooks()
+        delattr(self.extractor, "module_names")
         delattr(self.extractor, "output_type")

From 8827ec655e3affd76657ec85dcf19971184a2e03 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Tue, 29 Jul 2025 14:15:50 +0200
Subject: [PATCH 02/21] added option to specify list of module names instead of
 single name

---
 thingsvision/core/extraction/tensorflow.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/thingsvision/core/extraction/tensorflow.py b/thingsvision/core/extraction/tensorflow.py
index 31957710..86393e71 100644
--- a/thingsvision/core/extraction/tensorflow.py
+++ b/thingsvision/core/extraction/tensorflow.py
@@ -40,7 +40,7 @@ def _extract_batch(
         batch: Array,
         module_name: Optional[str] = None,
         module_names: Optional[List[str]] = None,
-        flatten_acts: bool,
+        flatten_acts: bool = False,
     ) -> Array:
         assert (
             module_name ^ module_names

From 60172b1d462237f14fe8ae0cd8492e6073e73e80 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Tue, 29 Jul 2025 14:17:29 +0200
Subject: [PATCH 03/21] added option to specify list of module names instead of
 single name

---
 thingsvision/core/extraction/base.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index 4c67409b..df8190f9 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -78,8 +78,8 @@ def extract_batch(
         batch: Union[TensorType["b", "c", "h", "w"], Array],
         module_name: Optional[str] = None,
         module_names: Optional[List[str]] = None,
-        flatten_acts: bool,
-        output_type: str,
+        flatten_acts: bool = False,
+        output_type: str = "ndarray",
     ) -> Dict[
         str,
         Union[
@@ -118,7 +118,7 @@ def _extract_batch(
         batch: Union[TensorType["b", "c", "h", "w"], Array],
         module_name: Optional[str] = None,
         module_names: Optional[List[str]] = None,
-        flatten_acts: bool,
+        flatten_acts: bool = False,
     ) -> Dict[
         str,
         Union[

From a195ed3a9b8548c4c76de1cab26aaa6a65b7f6d0 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Tue, 29 Jul 2025 15:03:27 +0200
Subject: [PATCH 04/21] added option to specify list of module names instead of
 single name

---
 thingsvision/core/extraction/base.py       | 2 +-
 thingsvision/core/extraction/tensorflow.py | 4 ++--
 thingsvision/core/extraction/torch.py      | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index df8190f9..cd318a45 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -214,7 +214,7 @@ def extract_features(
             Returns the feature matrix (e.g., $X \in \mathbb{R}^{n \times d}$ if penultimate or logits layer or flatten_acts = True).
         """
         assert (
-            module_name ^ module_names
+            bool(module_name) ^ bool(module_names)
         ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
         if module_name:
             module_names = [module_name]
diff --git a/thingsvision/core/extraction/tensorflow.py b/thingsvision/core/extraction/tensorflow.py
index 86393e71..a3470710 100644
--- a/thingsvision/core/extraction/tensorflow.py
+++ b/thingsvision/core/extraction/tensorflow.py
@@ -43,7 +43,7 @@ def _extract_batch(
         flatten_acts: bool = False,
     ) -> Array:
         assert (
-            module_name ^ module_names
+            bool(module_name) ^ bool(module_names)
         ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
         if module_name:
             module_names = [module_name]
@@ -66,7 +66,7 @@ def extract_batch(
         output_type: str = "ndarray",
     ) -> Array:
         assert (
-            module_name ^ module_names
+            bool(module_name) ^ bool(module_names)
         ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
         if module_name:
             module_names = [module_name]
diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
index 3dcf34f2..63d5eb73 100644
--- a/thingsvision/core/extraction/torch.py
+++ b/thingsvision/core/extraction/torch.py
@@ -110,7 +110,7 @@ def extract_batch(
         TensorType["b", "d"],
     ]:
         act = self._extract_batch(
-            batch=batch, module_name=self.module_names, flatten_acts=flatten_acts
+            batch=batch, module_names=self.module_names, flatten_acts=flatten_acts
         )
         if self.output_type == "ndarray":
             act = self._to_numpy(act)

From b1b3bdccb70fd4b4e007c96ee9ea0d441530b44c Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Tue, 29 Jul 2025 15:22:23 +0200
Subject: [PATCH 05/21] added option to specify list of module names instead of
 single name

---
 thingsvision/core/extraction/torch.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
index 63d5eb73..fd724735 100644
--- a/thingsvision/core/extraction/torch.py
+++ b/thingsvision/core/extraction/torch.py
@@ -174,6 +174,9 @@ def extract_features(
         output_dir: Optional[str] = None,
         step_size: Optional[int] = None,
     ):
+        assert (
+            bool(module_name) ^ bool(module_names)
+        ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
         self.model = self.model.to(self.device)
         self.activations = {}
         if module_name:
@@ -181,7 +184,6 @@ def extract_features(
         self._register_hooks(module_names=module_names)
         features = super().extract_features(
             batches=batches,
-            module_name=module_name,
             module_names=module_names,
             flatten_acts=flatten_acts,
             output_type=output_type,

From eb7755822cd606bb97ca3dcfb4db18b19d4127e0 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Tue, 29 Jul 2025 15:25:54 +0200
Subject: [PATCH 06/21] added option to specify list of module names instead of
 single name

---
 thingsvision/core/extraction/base.py | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index cd318a45..c00dd519 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -286,18 +286,15 @@ def extract_features(
 
     @staticmethod
     def _to_numpy(
-        features: Dict[
-            str,
-            Union[
-                TensorType["n", "num_maps", "h_prime", "w_prime"],
-                TensorType["n", "t", "d"],
-                TensorType["n", "p"],
-                TensorType["n", "d"],
-            ],
-        ],
-    ) -> Dict[str, Array]:
+        features: Union[
+            TensorType["n", "num_maps", "h_prime", "w_prime"],
+            TensorType["n", "t", "d"],
+            TensorType["n", "p"],
+            TensorType["n", "d"],
+        ]
+    ) -> Array:
         """Move activations to CPU and convert torch.Tensor to np.ndarray."""
-        return {k: v.cpu().numpy() for k, v in features.items()}
+        return features.numpy()
 
     def get_transformations(
         self, resize_dim: int = 256, crop_dim: int = 224, apply_center_crop: bool = True

From 1f143ff4d948430c5d2b7b90d0310706b7a4f126 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Tue, 29 Jul 2025 15:28:43 +0200
Subject: [PATCH 07/21] added option to specify list of module names instead of
 single name

---
 thingsvision/core/extraction/base.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index c00dd519..5dd58495 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -276,12 +276,12 @@ def extract_features(
         else:
             for module_name in module_names:
                 if self.get_backend() == "pt":
-                    features = torch.cat(features[module_name])
+                    features[module_name] = torch.cat(features[module_name])
                     if output_type == "ndarray":
-                        features = self._to_numpy(features)
+                        features[module_name] = self._to_numpy(features[module_name])
                 else:
-                    features = np.vstack(features[module_name])
-                    print(f"...Features shape: {features.shape}")
+                    features[module_name] = np.vstack(features[module_name])
+                    print(f"...Features shape: {features[module_name].shape}")
         return features
 
     @staticmethod

From fa3df920f1e6f55ffc562bc845eb80e13de637a6 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Mon, 4 Aug 2025 11:45:27 +0200
Subject: [PATCH 08/21] adjustments for backward compatibility

---
 tests/test_features.py                     | 50 ++++++++++++++++++++++
 thingsvision/core/extraction/base.py       |  3 ++
 thingsvision/core/extraction/tensorflow.py | 30 ++++---------
 thingsvision/core/extraction/torch.py      |  4 +-
 4 files changed, 63 insertions(+), 24 deletions(-)

diff --git a/tests/test_features.py b/tests/test_features.py
index 95334626..3c92bdc2 100644
--- a/tests/test_features.py
+++ b/tests/test_features.py
@@ -44,6 +44,19 @@ def get_4D_features(self):
             flatten_acts=False,
         )
         return features
+    
+    def get_multi_features(self):
+        model_name = "vgg16_bn"
+        extractor, _, batches = helper.create_extractor_and_dataloader(
+            model_name=model_name, pretrained=False, source="torchvision"
+        )
+        module_names = ["features.23", "classifier.3"]
+        features = extractor.extract_features(
+            batches=batches,
+            module_names=module_names,
+            flatten_acts=False,
+        )
+        return features
 
     def test_postprocessing(self):
         """Test different postprocessing methods (e.g., centering, normalization, compression)."""
@@ -89,6 +102,18 @@ def test_storing_4d(self):
             )
 
             self.check_file_exists("features", format, False)
+    
+    def test_storing_multi(self):
+        features = self.get_multi_features()
+        for _, feature in features.items():
+            for format in helper.FILE_FORMATS:
+                # tests whether features can be saved in any of the formats
+                save_features(
+                    features=feature,
+                    out_path=f"{helper.OUT_PATH}",
+                    file_format=format,
+                )
+                self.check_file_exists(f"features", format, False)
 
     def test_splitting_2d(self):
         n_splits = 3
@@ -129,3 +154,28 @@ def test_splitting_4d(self):
                 file_format="txt",
                 n_splits=n_splits,
             )
+    
+    def test_splitting_multi(self):
+        n_splits = 3
+        features = self.get_multi_features()
+        for _, feature in features.items():
+            for format in set(helper.FILE_FORMATS) - set(["txt"]):
+                if format == "pt":
+                    feature = torch.from_numpy(feature)
+                split_features(
+                    features=feature,
+                    root=helper.OUT_PATH,
+                    file_format=format,
+                    n_splits=n_splits,
+                )
+
+                for i in range(1, n_splits):
+                    self.check_file_exists(f"features_{i:02d}", format, False)
+
+            with self.assertRaises(Exception):
+                split_features(
+                    features=feature,
+                    root=helper.OUT_PATH,
+                    file_format="txt",
+                    n_splits=n_splits,
+                )
diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index 5dd58495..5fdc36b1 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -282,6 +282,9 @@ def extract_features(
                 else:
                     features[module_name] = np.vstack(features[module_name])
                     print(f"...Features shape: {features[module_name].shape}")
+        if module_name is not None:
+            # for backward compatibility
+            features = features[module_name]
         return features
 
     @staticmethod
diff --git a/thingsvision/core/extraction/tensorflow.py b/thingsvision/core/extraction/tensorflow.py
index a3470710..9a30befd 100644
--- a/thingsvision/core/extraction/tensorflow.py
+++ b/thingsvision/core/extraction/tensorflow.py
@@ -38,19 +38,15 @@ def __init__(
     def _extract_batch(
         self,
         batch: Array,
-        module_name: Optional[str] = None,
+        module_name: str,
         module_names: Optional[List[str]] = None,
         flatten_acts: bool = False,
     ) -> Array:
-        assert (
-            bool(module_name) ^ bool(module_names)
-        ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
-        if module_name:
-            module_names = [module_name]
-        layer_outs = [self.model.get_layer(name).output for name in module_names]
+        assert module_names is None, "TensorFlowExtractor does not support multiple module names."
+        layer_out = [self.model.get_layer(module_name).output]
         activation_model = keras.models.Model(
             inputs=self.model.input,
-            outputs=layer_outs,
+            outputs=layer_out,
         )
         activations = activation_model.predict(batch)
         if flatten_acts:
@@ -60,22 +56,14 @@ def _extract_batch(
     def extract_batch(
         self,
         batch: Array,
-        module_name: Optional[str] = None,
+        module_name: str,
         module_names: Optional[List[str]] = None,
         flatten_acts: bool = False,
         output_type: str = "ndarray",
     ) -> Array:
-        assert (
-            bool(module_name) ^ bool(module_names)
-        ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
-        if module_name:
-            module_names = [module_name]
-        self.model = self.model.to(self.device)
-        self.activations = {}
-        if module_name:
-            module_names = [module_name]
-        self._module_and_output_check(module_names, output_type)
-        activations = self._extract_batch(batch, module_names, flatten_acts)
+        assert module_names is None, "TensorFlowExtractor does not support multiple module names."
+        self._module_and_output_check(module_name, output_type)
+        activations = self._extract_batch(batch, module_name, flatten_acts)
         return activations
 
     def show_model(self) -> str:
@@ -123,4 +111,4 @@ def get_default_transformation(
         return composition
 
     def get_backend(self) -> str:
-        return "tf"
+        return "tf"
\ No newline at end of file
diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
index fd724735..229889de 100644
--- a/thingsvision/core/extraction/torch.py
+++ b/thingsvision/core/extraction/torch.py
@@ -190,9 +190,7 @@ def extract_features(
             output_dir=output_dir,
             step_size=step_size,
         )
-        if self.hook_handles:
-            for handle in self.hook_handles:
-                handle.remove()
+        self._unregister_hooks()
         return features
 
     def forward(

From d603a32f7043e2ee486ec1d14a6887f877ddbdbc Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Mon, 4 Aug 2025 11:53:49 +0200
Subject: [PATCH 09/21] adressing comments

---
 thingsvision/core/extraction/base.py       | 16 ++++++++++------
 thingsvision/core/extraction/tensorflow.py | 10 +++++++---
 thingsvision/core/extraction/torch.py      |  7 ++++---
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index 5fdc36b1..3c85eddd 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -3,6 +3,7 @@
 import re
 import warnings
 from typing import Callable, Dict, Iterator, List, Optional, Union
+from collections import defaultdict
 
 import numpy as np
 from torchtyping import TensorType
@@ -98,6 +99,8 @@ def extract_batch(
         ----------
         batch : np.ndarray or torch.Tensor
             mini-batch of three-dimensional image tensors.
+        module_name : str
+            Name of the neural network layer for which features should be extracted.
         module_names : List[str]
             Names of the modules for which features should be extracted.
         flatten_acts : bool
@@ -213,9 +216,10 @@ def extract_features(
         output : np.ndarray or torch.Tensor
             Returns the feature matrix (e.g., $X \in \mathbb{R}^{n \times d}$ if penultimate or logits layer or flatten_acts = True).
         """
-        assert (
-            bool(module_name) ^ bool(module_names)
-        ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
+        if not bool(module_name) ^ bool(module_names):
+            raise ValueError(
+                "\nPlease provide either a single module name or a list of module names, but not both.\n"
+            )
         if module_name:
             module_names = [module_name]
         self._module_and_output_check(module_names, output_type)
@@ -227,7 +231,7 @@ def extract_features(
                 step_size = 8000 // (len(next(iter(batches))) * 3) + 1
 
         # create feature dict per module name
-        features = {name: [] for name in module_names}
+        features = defaultdict(list)
         image_ct, last_image_ct = 0, 0
         for i, batch in tqdm(
             enumerate(batches, start=1), desc="Batch", total=len(batches)
@@ -265,7 +269,7 @@ def extract_features(
                         )
                         features_subset = np.vstack(features[module_name])
                         np.save(features_subset_file, features_subset)
-                    features = {name: [] for name in module_names}
+                    features = defaultdict(list)
                     last_image_ct = image_ct
         print(
             f"...Features successfully extracted for all {image_ct} images in the database."
@@ -294,7 +298,7 @@ def _to_numpy(
             TensorType["n", "t", "d"],
             TensorType["n", "p"],
             TensorType["n", "d"],
-        ]
+        ],
     ) -> Array:
         """Move activations to CPU and convert torch.Tensor to np.ndarray."""
         return features.numpy()
diff --git a/thingsvision/core/extraction/tensorflow.py b/thingsvision/core/extraction/tensorflow.py
index 9a30befd..dca9ab2b 100644
--- a/thingsvision/core/extraction/tensorflow.py
+++ b/thingsvision/core/extraction/tensorflow.py
@@ -42,7 +42,9 @@ def _extract_batch(
         module_names: Optional[List[str]] = None,
         flatten_acts: bool = False,
     ) -> Array:
-        assert module_names is None, "TensorFlowExtractor does not support multiple module names."
+        assert (
+            module_names is None
+        ), "TensorFlowExtractor does not support multiple module names."
         layer_out = [self.model.get_layer(module_name).output]
         activation_model = keras.models.Model(
             inputs=self.model.input,
@@ -61,7 +63,9 @@ def extract_batch(
         flatten_acts: bool = False,
         output_type: str = "ndarray",
     ) -> Array:
-        assert module_names is None, "TensorFlowExtractor does not support multiple module names."
+        assert (
+            module_names is None
+        ), "TensorFlowExtractor does not support multiple module names."
         self._module_and_output_check(module_name, output_type)
         activations = self._extract_batch(batch, module_name, flatten_acts)
         return activations
@@ -111,4 +115,4 @@ def get_default_transformation(
         return composition
 
     def get_backend(self) -> str:
-        return "tf"
\ No newline at end of file
+        return "tf"
diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
index 229889de..5644021e 100644
--- a/thingsvision/core/extraction/torch.py
+++ b/thingsvision/core/extraction/torch.py
@@ -174,9 +174,10 @@ def extract_features(
         output_dir: Optional[str] = None,
         step_size: Optional[int] = None,
     ):
-        assert (
-            bool(module_name) ^ bool(module_names)
-        ), "Please provide either a single module name or a list of module names for which features should be extracted.\n"
+        if not bool(module_name) ^ bool(module_names):
+            raise ValueError(
+                "\nPlease provide either a single module name or a list of module names, but not both.\n"
+            )
         self.model = self.model.to(self.device)
         self.activations = {}
         if module_name:

From a0750242b99aad25fb9a848280c6b718cb150d88 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Mon, 4 Aug 2025 14:04:59 +0200
Subject: [PATCH 10/21] fixing tests and unifying interface across frameworks

---
 tests/test_features.py                     |  6 +--
 thingsvision/core/extraction/base.py       | 55 +++++++++++++++-------
 thingsvision/core/extraction/tensorflow.py | 46 ++++++++++--------
 thingsvision/core/extraction/torch.py      | 55 ++++++++++++++++++----
 4 files changed, 114 insertions(+), 48 deletions(-)

diff --git a/tests/test_features.py b/tests/test_features.py
index 3c92bdc2..6044389b 100644
--- a/tests/test_features.py
+++ b/tests/test_features.py
@@ -106,11 +106,11 @@ def test_storing_4d(self):
     def test_storing_multi(self):
         features = self.get_multi_features()
         for _, feature in features.items():
-            for format in helper.FILE_FORMATS:
-                # tests whether features can be saved in any of the formats
+            for format in set(helper.FILE_FORMATS) - set(["txt"]):
+                # tests whether features can be saved in any of the formats except txt
                 save_features(
                     features=feature,
-                    out_path=f"{helper.OUT_PATH}",
+                    out_path=helper.OUT_PATH,
                     file_format=format,
                 )
                 self.check_file_exists(f"features", format, False)
diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index 3c85eddd..458c0244 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -81,15 +81,26 @@ def extract_batch(
         module_names: Optional[List[str]] = None,
         flatten_acts: bool = False,
         output_type: str = "ndarray",
-    ) -> Dict[
-        str,
-        Union[
+    ) -> Union[
+        # This is the return type when 'module_names' is used
+        Dict[
+            str,
             Union[
-                TensorType["n", "num_maps", "h_prime", "w_prime"],
-                TensorType["n", "t", "d"],
-                TensorType["n", "p"],
-                TensorType["n", "d"],
+                Union[
+                    TensorType["n", "num_maps", "h_prime", "w_prime"],
+                    TensorType["n", "t", "d"],
+                    TensorType["n", "p"],
+                    TensorType["n", "d"],
+                ],
+                Array,
             ],
+        ],
+        # This is the return type when 'module_name' is used (for backward compatibility)
+        Union[
+            TensorType["n", "num_maps", "h_prime", "w_prime"],
+            TensorType["n", "t", "d"],
+            TensorType["n", "p"],
+            TensorType["n", "d"],
             Array,
         ],
     ]:
@@ -119,8 +130,7 @@ def extract_batch(
     def _extract_batch(
         self,
         batch: Union[TensorType["b", "c", "h", "w"], Array],
-        module_name: Optional[str] = None,
-        module_names: Optional[List[str]] = None,
+        module_names: List[str],
         flatten_acts: bool = False,
     ) -> Dict[
         str,
@@ -163,15 +173,26 @@ def extract_features(
         output_type: Optional[str] = "ndarray",
         output_dir: Optional[str] = None,
         step_size: Optional[int] = None,
-    ) -> Dict[
-        str,
-        Union[
+    ) -> Union[
+        # This is the return type when 'module_names' is used
+        Dict[
+            str,
             Union[
-                TensorType["n", "num_maps", "h_prime", "w_prime"],
-                TensorType["n", "t", "d"],
-                TensorType["n", "p"],
-                TensorType["n", "d"],
+                Union[
+                    TensorType["n", "num_maps", "h_prime", "w_prime"],
+                    TensorType["n", "t", "d"],
+                    TensorType["n", "p"],
+                    TensorType["n", "d"],
+                ],
+                Array,
             ],
+        ],
+        # This is the return type when 'module_name' is used (for backward compatibility)
+        Union[
+            TensorType["n", "num_maps", "h_prime", "w_prime"],
+            TensorType["n", "t", "d"],
+            TensorType["n", "p"],
+            TensorType["n", "d"],
             Array,
         ],
     ]:
@@ -288,7 +309,7 @@ def extract_features(
                     print(f"...Features shape: {features[module_name].shape}")
         if module_name is not None:
             # for backward compatibility
-            features = features[module_name]
+            return features[module_name]
         return features
 
     @staticmethod
diff --git a/thingsvision/core/extraction/tensorflow.py b/thingsvision/core/extraction/tensorflow.py
index dca9ab2b..aff33201 100644
--- a/thingsvision/core/extraction/tensorflow.py
+++ b/thingsvision/core/extraction/tensorflow.py
@@ -38,36 +38,44 @@ def __init__(
     def _extract_batch(
         self,
         batch: Array,
-        module_name: str,
-        module_names: Optional[List[str]] = None,
-        flatten_acts: bool = False,
-    ) -> Array:
-        assert (
-            module_names is None
-        ), "TensorFlowExtractor does not support multiple module names."
-        layer_out = [self.model.get_layer(module_name).output]
+        module_names: Optional[List[str]],
+        flatten_acts: bool,
+    ) -> Dict[str, Array]:
+        layer_outputs = [self.model.get_layer(name).output for name in module_names]
         activation_model = keras.models.Model(
             inputs=self.model.input,
-            outputs=layer_out,
+            outputs=layer_outputs,
         )
-        activations = activation_model.predict(batch)
+        activations_list = activation_model.predict(batch)
+        if len(module_names) == 1:
+            activations_list = [activations_list]
+        activations_dict = {
+            name: act for name, act in zip(module_names, activations_list)
+        }
         if flatten_acts:
-            activations = activations.reshape(activations.shape[0], -1)
-        return activations
+            for name, act in activations_dict.items():
+                activations_dict[name] = act.reshape(act.shape[0], -1)
+        return activations_dict
 
     def extract_batch(
         self,
         batch: Array,
-        module_name: str,
+        module_name: Optional[str] = None,
         module_names: Optional[List[str]] = None,
         flatten_acts: bool = False,
         output_type: str = "ndarray",
-    ) -> Array:
-        assert (
-            module_names is None
-        ), "TensorFlowExtractor does not support multiple module names."
-        self._module_and_output_check(module_name, output_type)
-        activations = self._extract_batch(batch, module_name, flatten_acts)
+    ) -> Union[Array, Dict[str, Array]]:
+        if not bool(module_name) ^ bool(module_names):
+            raise ValueError(
+                "\nPlease provide either a single module name or a list of module names, but not both.\n"
+            )
+        if not module_names:
+            module_names = [module_name]
+        self._module_and_output_check(module_names, output_type)
+        # Extract features from the specified module, tensorflow does not support multiple modules extraction
+        activations = self._extract_batch(batch, module_names, flatten_acts)
+        if module_name:
+            return activations[module_name]
         return activations
 
     def show_model(self) -> str:
diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
index 5644021e..84ed1360 100644
--- a/thingsvision/core/extraction/torch.py
+++ b/thingsvision/core/extraction/torch.py
@@ -104,17 +104,40 @@ def extract_batch(
         batch: TensorType["b", "c", "h", "w"],
         flatten_acts: bool,
     ) -> Union[
-        TensorType["b", "num_maps", "h_prime", "w_prime"],
-        TensorType["b", "t", "d"],
-        TensorType["b", "p"],
-        TensorType["b", "d"],
+        # This is the return type when 'module_names' is used
+        Dict[
+            str,
+            Union[
+                Union[
+                    TensorType["n", "num_maps", "h_prime", "w_prime"],
+                    TensorType["n", "t", "d"],
+                    TensorType["n", "p"],
+                    TensorType["n", "d"],
+                ],
+                Array,
+            ],
+        ],
+        # This is the return type when 'module_name' is used (for backward compatibility)
+        Union[
+            TensorType["n", "num_maps", "h_prime", "w_prime"],
+            TensorType["n", "t", "d"],
+            TensorType["n", "p"],
+            TensorType["n", "d"],
+            Array,
+        ],
     ]:
-        act = self._extract_batch(
-            batch=batch, module_names=self.module_names, flatten_acts=flatten_acts
+        acts = self._extract_batch(
+            batch=batch,
+            module_name=getattr(self, "module_name", None),
+            module_names=getattr(self, "module_names", None),
+            flatten_acts=flatten_acts,
         )
         if self.output_type == "ndarray":
-            act = self._to_numpy(act)
-        return act
+            for module_name, act in acts.items():
+                acts[module_name] = self._to_numpy(act)
+        if getattr(self, "module_name", None) is not None:
+            return acts[self.module_name]
+        return acts
 
     @torch.no_grad()
     def _extract_batch(
@@ -291,7 +314,11 @@ def get_backend(self) -> str:
 class BatchExtraction(object):
 
     def __init__(
-        self, extractor: PyTorchExtractor, module_names: List[str], output_type: str
+        self,
+        extractor: PyTorchExtractor,
+        module_name: Optional[List[str]] = None,
+        module_names: Optional[List[str]] = None,
+        output_type: str = "ndarray",
     ) -> None:
         """
         Mini-batch extraction object that can be used as a with-statement in a PyTorch extractor.
@@ -300,10 +327,18 @@ def __init__(
         ----------
         extractor (object): PyTorchExtractor class.
         module_name (str): The module of model for which features will be extracted.
+        module_names (List[str]): List of modules of model for which features will be extracted.
         output_type (str): Type of the feature matrix returned by the extractor.
 
         """
+        if not bool(module_name) ^ bool(module_names):
+            raise ValueError(
+                "\nPlease provide either a single module name or a list of module names, but not both.\n"
+            )
+        if module_name:
+            module_names = [module_name]
         self.extractor = extractor
+        self.module_name = module_name
         self.module_names = module_names
         self.output_type = output_type
 
@@ -311,6 +346,7 @@ def __enter__(self) -> PyTorchExtractor:
         """Registering hooks and setting attributes during opening."""
         self.extractor._module_and_output_check(self.module_names, self.output_type)
         self.extractor._register_hooks(self.module_names)
+        setattr(self.extractor, "module_name", self.module_name)
         setattr(self.extractor, "module_names", self.module_names)
         setattr(self.extractor, "output_type", self.output_type)
         return self.extractor
@@ -318,5 +354,6 @@ def __enter__(self) -> PyTorchExtractor:
     def __exit__(self, *args):
         """Removing hooks and deleting attributes at closing."""
         self.extractor._unregister_hooks()
+        delattr(self.extractor, "module_name")
         delattr(self.extractor, "module_names")
         delattr(self.extractor, "output_type")

From c6e3902eba2313da8535729e218aae99b1ee167a Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Mon, 4 Aug 2025 14:51:30 +0200
Subject: [PATCH 11/21] fixing tests

---
 tests/test_features.py                |  7 ++++---
 thingsvision/core/extraction/base.py  |  7 +++++--
 thingsvision/core/extraction/torch.py | 26 +++++++++++++++++---------
 3 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/tests/test_features.py b/tests/test_features.py
index 6044389b..0dbd2668 100644
--- a/tests/test_features.py
+++ b/tests/test_features.py
@@ -158,8 +158,8 @@ def test_splitting_4d(self):
     def test_splitting_multi(self):
         n_splits = 3
         features = self.get_multi_features()
-        for _, feature in features.items():
-            for format in set(helper.FILE_FORMATS) - set(["txt"]):
+        for format in set(helper.FILE_FORMATS) - set(["txt"]):
+            for _, feature in features.items():
                 if format == "pt":
                     feature = torch.from_numpy(feature)
                 split_features(
@@ -172,7 +172,8 @@ def test_splitting_multi(self):
                 for i in range(1, n_splits):
                     self.check_file_exists(f"features_{i:02d}", format, False)
 
-            with self.assertRaises(Exception):
+        with self.assertRaises(Exception):
+            for _, feature in features.items():
                 split_features(
                     features=feature,
                     root=helper.OUT_PATH,
diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index 458c0244..4bbdd6ad 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -241,8 +241,11 @@ def extract_features(
             raise ValueError(
                 "\nPlease provide either a single module name or a list of module names, but not both.\n"
             )
-        if module_name:
+        if module_name is not None:
+            single_module_call = True
             module_names = [module_name]
+        else:
+            single_module_call = False
         self._module_and_output_check(module_names, output_type)
         if output_dir:
             os.makedirs(output_dir, exist_ok=True)
@@ -307,7 +310,7 @@ def extract_features(
                 else:
                     features[module_name] = np.vstack(features[module_name])
                     print(f"...Features shape: {features[module_name].shape}")
-        if module_name is not None:
+        if single_module_call:
             # for backward compatibility
             return features[module_name]
         return features
diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
index 84ed1360..cb1dc748 100644
--- a/thingsvision/core/extraction/torch.py
+++ b/thingsvision/core/extraction/torch.py
@@ -93,10 +93,15 @@ def _unregister_hooks(self) -> None:
                 category=UserWarning,
             )
 
-    def batch_extraction(self, module_names: List[str], output_type: str) -> object:
+    def batch_extraction(
+        self,
+        module_name: Optional[str] = None,
+        module_names: Optional[List[str]] = None,
+        output_type: str = "ndarray",
+    ) -> object:
         """Allows mini-batch extraction for custom data pipeline using a with-statement."""
         return BatchExtraction(
-            extractor=self, module_names=module_names, output_type=output_type
+            extractor=self, module_name=module_name, module_names=module_names, output_type=output_type
         )
 
     def extract_batch(
@@ -128,8 +133,7 @@ def extract_batch(
     ]:
         acts = self._extract_batch(
             batch=batch,
-            module_name=getattr(self, "module_name", None),
-            module_names=getattr(self, "module_names", None),
+            module_names=self.module_names,
             flatten_acts=flatten_acts,
         )
         if self.output_type == "ndarray":
@@ -203,11 +207,13 @@ def extract_features(
             )
         self.model = self.model.to(self.device)
         self.activations = {}
-        if module_name:
-            module_names = [module_name]
-        self._register_hooks(module_names=module_names)
+        if module_name is not None:
+            self._register_hooks(module_names=[module_name])
+        else:
+            self._register_hooks(module_names=module_names)
         features = super().extract_features(
             batches=batches,
+            module_name=module_name,
             module_names=module_names,
             flatten_acts=flatten_acts,
             output_type=output_type,
@@ -316,7 +322,7 @@ class BatchExtraction(object):
     def __init__(
         self,
         extractor: PyTorchExtractor,
-        module_name: Optional[List[str]] = None,
+        module_name: Optional[str] = None,
         module_names: Optional[List[str]] = None,
         output_type: str = "ndarray",
     ) -> None:
@@ -335,7 +341,8 @@ def __init__(
             raise ValueError(
                 "\nPlease provide either a single module name or a list of module names, but not both.\n"
             )
-        if module_name:
+        print(module_names, module_name, "INITIAKLAA")
+        if module_name is not None:
             module_names = [module_name]
         self.extractor = extractor
         self.module_name = module_name
@@ -344,6 +351,7 @@ def __init__(
 
     def __enter__(self) -> PyTorchExtractor:
         """Registering hooks and setting attributes during opening."""
+        print("EXTRACTOR", self.module_names)
         self.extractor._module_and_output_check(self.module_names, self.output_type)
         self.extractor._register_hooks(self.module_names)
         setattr(self.extractor, "module_name", self.module_name)

From 21fae0b1b91b67ab467539ee0567a77fb37c0018 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Mon, 4 Aug 2025 15:02:23 +0200
Subject: [PATCH 12/21] update python

---
 .github/workflows/coverage.yml | 4 ++--
 .github/workflows/tests.yml    | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml
index a526f12f..6630f580 100644
--- a/.github/workflows/coverage.yml
+++ b/.github/workflows/coverage.yml
@@ -12,14 +12,14 @@ jobs:
         os: [ubuntu-latest]
     env:
       OS: ${{ matrix.os }}
-      PYTHON: '3.9'
+      PYTHON: '3.10'
 
     steps:
     - uses: actions/checkout@master
     - name: Setup Python
       uses: actions/setup-python@master
       with:
-        python-version: 3.9
+        python-version: 3.10
         
     - name: Install dependencies
       run: |
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 0d1cb095..1192c245 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.9]
+        python-version: ["3.10.x"]
 
     steps:
     - uses: actions/checkout@v2

From dde828f7b4125fb2f4110cfc12262c9a026da303 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Tue, 5 Aug 2025 15:04:08 +0200
Subject: [PATCH 13/21] merge upstream and integrate comments

---
 .github/workflows/coverage.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml
index 6630f580..d9f07d12 100644
--- a/.github/workflows/coverage.yml
+++ b/.github/workflows/coverage.yml
@@ -19,7 +19,7 @@ jobs:
     - name: Setup Python
       uses: actions/setup-python@master
       with:
-        python-version: 3.10
+        python-version: "3.10"
         
     - name: Install dependencies
       run: |

From 63c951d1414d0733c8352cb8190b3b825ed5cd79 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Tue, 5 Aug 2025 17:46:34 +0200
Subject: [PATCH 14/21] update readme, and added concatenate option for
 extract_features

---
 README.md                            | 42 ++++++++++++++++++++++++++++
 thingsvision/core/extraction/base.py | 34 ++++++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/README.md b/README.md
index cbfaa2b7..53be6fe4 100644
--- a/README.md
+++ b/README.md
@@ -252,6 +252,48 @@ for batch in my_dataloader:
   ... # whatever post-processing you want to add to the extracted features
 ```
 
+#### Multi Module Feature Extraction
+
+We've added the functionality to also jointly extract features of multiple `module_names`.
+
+##### PyTorch
+
+```python
+module_names = ['visual', ...] # add more module_names here
+
+# your custom dataset and dataloader classes come here (for example, a PyTorch data loader)
+my_dataset = ...
+my_dataloader = ...
+
+with extractor.batch_extraction(module_names=module_names, output_type="tensor") as e: 
+  for batch in my_dataloader:
+    ... # whatever preprocessing you want to add to the batch
+    feature_batch_dict = e.extract_batch(
+      batch=batch,
+      flatten_acts=True, # flatten 2D feature maps from an early convolutional or attention layer
+      )
+    ... # whatever post-processing you want to add to the extracted features
+```
+
+##### TensorFlow / Keras
+
+```python
+module_names = ['visual', ...] # add more module_names here
+
+# your custom dataset and dataloader classes come here (for example, TFRecords files)
+my_dataset = ...
+my_dataloader = ...
+
+for batch in my_dataloader:
+  ... # whatever preprocessing you want to add to the batch
+  feature_batch = extractor.extract_batch(
+    batch=batch,
+    module_names=module_names,
+    flatten_acts=True, # flatten 2D feature maps from an early convolutional or attention layer
+    )
+  ... # whatever post-processing you want to add to the extracted features
+```
+
 #### Human alignment
 
 *Human alignment*: If you want to align the extracted features with human object similarity according to the approach introduced in *[Improving neural network representations using human similiarty judgments](https://proceedings.neurips.cc/paper_files/paper/2023/hash/9febda1c8344cc5f2d51713964864e93-Abstract-Conference.html)* you can optionally `align` the extracted features using the following method:
diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index a78d4630..64348121 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -174,6 +174,7 @@ def extract_features(
         output_dir: Optional[str] = None,
         step_size: Optional[int] = None,
         file_name_suffix: str = "",
+        save_in_one_file: bool = False,
     ) -> Union[
         # This is the return type when 'module_names' is used
         Dict[
@@ -234,6 +235,11 @@ def extract_features(
             Only used if output_dir is defined.
         file_name_suffix: str
             Suffix to append to the output file names (e.g., "_train", "_val").
+        save_in_one_file : bool
+            If True, all features are saved in one file. If output_dir is defined,
+            the features are saved in separate files for each module name. They are first
+            saved in chunks of step_size batches, and then all features are concatenated
+            and saved in one file.
         Returns
         -------
         output : np.ndarray or torch.Tensor
@@ -301,6 +307,34 @@ def extract_features(
             f"...Features successfully extracted for all {image_ct} images in the database."
         )
         if output_dir:
+            if save_in_one_file:
+                # load features per module name and concatenate them
+                for module_name in module_names:
+                    # load from files
+                    features = []
+                    for file in sorted(
+                        os.listdir(os.path.join(output_dir, module_name))
+                    ):
+                        if file.endswith(".npy") or file.endswith(".pt"):
+                            features.append(
+                                np.load(os.path.join(output_dir, module_name, file))
+                                if file.endswith(".npy")
+                                else torch.load(
+                                    os.path.join(output_dir, module_name, file)
+                                )
+                            )
+                    features = np.concatenate(features) if output_type == "ndarray" else torch.cat(features)
+                    features_file = os.path.join(
+                        output_dir, f"{module_name}/features{file_name_suffix}"
+                    )
+                    if output_type == "ndarray":
+                        np.save(f"{features_file}.npy", features)
+                    else:  # output_type = tensor
+                        torch.save(features, f"{features_file}.pt")
+                    print(
+                        f"...Features for module '{module_name}' were saved to {features_file}."
+                    )
+                
             print(f"...Features were saved to {output_dir}.")
             return None
         else:

From 6408a2882f842786fc0a83a899d7e221d51badfa Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Tue, 5 Aug 2025 17:51:01 +0200
Subject: [PATCH 15/21] update readme, and added concatenate option for
 extract_features

---
 thingsvision/core/extraction/base.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index 64348121..d8e42435 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -315,22 +315,23 @@ def extract_features(
                     for file in sorted(
                         os.listdir(os.path.join(output_dir, module_name))
                     ):
-                        if file.endswith(".npy") or file.endswith(".pt"):
-                            features.append(
-                                np.load(os.path.join(output_dir, module_name, file))
-                                if file.endswith(".npy")
-                                else torch.load(
-                                    os.path.join(output_dir, module_name, file)
+                        if self.get_backend() == "pt" and output_type != "ndarray":
+                            if file.endswith(".pt"):
+                                features.append(
+                                    torch.load(os.path.join(output_dir, module_name, file))
+                                )
+                        else:
+                            if file.endswith(".npy"):
+                                features.append(
+                                    np.load(os.path.join(output_dir, module_name, file))
                                 )
-                            )
-                    features = np.concatenate(features) if output_type == "ndarray" else torch.cat(features)
                     features_file = os.path.join(
                         output_dir, f"{module_name}/features{file_name_suffix}"
                     )
                     if output_type == "ndarray":
-                        np.save(f"{features_file}.npy", features)
+                        np.save(f"{features_file}.npy", np.concatenate(features))
                     else:  # output_type = tensor
-                        torch.save(features, f"{features_file}.pt")
+                        torch.save(torch.cat(features), f"{features_file}.pt")
                     print(
                         f"...Features for module '{module_name}' were saved to {features_file}."
                     )

From 34b135ca86e77b25453c486bb0e7c7085cd137d3 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Wed, 6 Aug 2025 11:48:59 +0200
Subject: [PATCH 16/21] update readme and concatenated saving

---
 README.md                            | 21 +++------------------
 thingsvision/core/extraction/base.py | 13 ++++++++-----
 2 files changed, 11 insertions(+), 23 deletions(-)

diff --git a/README.md b/README.md
index 53be6fe4..fc564c8b 100644
--- a/README.md
+++ b/README.md
@@ -97,32 +97,17 @@ Neural networks come from different sources. With `thingsvision`, you can extrac
 <!-- Setting up your environment -->
 ### :computer: Setting up your environment
 #### Working locally
-First, create a new `conda environment` with Python version 3.8, 3.9, 3.10, or 3.11 e.g. by using `conda`:
-
+First, create a new `conda environment` with Python version 3.10, or 3.11 e.g. by using `conda`:
 ```bash
-$ conda create -n thingsvision python=3.9
+$ conda create -n thingsvision python=3.10
 $ conda activate thingsvision
 ```
-
 Then, activate the environment and simply install `thingsvision` via running the following `pip` command in your terminal.
-
 ```bash
 $ pip install --upgrade thingsvision
-$ pip install git+https://github.com/openai/CLIP.git
-```
-
-If you want to extract features for [harmonized models](https://vicco-group.github.io/thingsvision/AvailableModels.html#harmonization) from the [Harmonization repo](https://github.com/serre-lab/harmonization), you have to additionally run the following `pip` command in your `thingsvision` environment (FYI: as of now, this seems to be working smoothly on Ubuntu only but not on macOS),
-
-```bash
-$ pip install git+https://github.com/serre-lab/Harmonization.git
-$ pip install keras-cv-attention-models>=1.3.5
 ```
 
-If you want to extract features for [DreamSim](https://dreamsim-nights.github.io/) from the [DreamSim repo](https://github.com/ssundaram21/dreamsim), you have to additionally run the following `pip` command in your `thingsvision` environment,
-
-```bash
-$ pip install dreamsim==0.1.2
-```
+The package automatically installs the [Harmonization](https://github.com/serre-lab/harmonization) and [DreamSim](https://github.com/ssundaram21/dreamsim) repositories. See the documentation for available [harmonized models](https://vicco-group.github.io/thingsvision/AvailableModels.html#harmonization) and [DreamSim models](https://vicco-group.github.io/thingsvision/AvailableModels.html#dreamsim) in `thingsvision`.
 
 See the [docs](https://vicco-group.github.io/thingsvision/AvailableModels.html#dreamsim) for which `DreamSim` models are available in `thingsvision`.
 
diff --git a/thingsvision/core/extraction/base.py b/thingsvision/core/extraction/base.py
index d8e42435..6b5db03c 100644
--- a/thingsvision/core/extraction/base.py
+++ b/thingsvision/core/extraction/base.py
@@ -264,6 +264,7 @@ def extract_features(
 
         # create feature dict per module name
         features = defaultdict(list)
+        feature_file_names = defaultdict(list)
         image_ct, last_image_ct = 0, 0
         for i, batch in tqdm(
             enumerate(batches, start=1), desc="Batch", total=len(batches)
@@ -303,6 +304,7 @@ def extract_features(
                         np.save(features_subset_file, features_subset)
                     features = defaultdict(list)
                     last_image_ct = image_ct
+                    feature_file_names[module_name].append(features_subset_file)
         print(
             f"...Features successfully extracted for all {image_ct} images in the database."
         )
@@ -312,18 +314,16 @@ def extract_features(
                 for module_name in module_names:
                     # load from files
                     features = []
-                    for file in sorted(
-                        os.listdir(os.path.join(output_dir, module_name))
-                    ):
+                    for file in feature_file_names[module_name]:
                         if self.get_backend() == "pt" and output_type != "ndarray":
                             if file.endswith(".pt"):
                                 features.append(
-                                    torch.load(os.path.join(output_dir, module_name, file))
+                                    torch.load(os.path.join(output_dir, file))
                                 )
                         else:
                             if file.endswith(".npy"):
                                 features.append(
-                                    np.load(os.path.join(output_dir, module_name, file))
+                                    np.load(os.path.join(output_dir, file))
                                 )
                     features_file = os.path.join(
                         output_dir, f"{module_name}/features{file_name_suffix}"
@@ -335,6 +335,9 @@ def extract_features(
                     print(
                         f"...Features for module '{module_name}' were saved to {features_file}."
                     )
+                    # remove temporary files
+                    for file in feature_file_names[module_name]:
+                        os.remove(os.path.join(output_dir, file))
                 
             print(f"...Features were saved to {output_dir}.")
             return None

From 3783ad22ff8882e4173f3380a8d5f9e8ce15dd76 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Wed, 6 Aug 2025 11:51:55 +0200
Subject: [PATCH 17/21] small bug fix

---
 thingsvision/core/extraction/torch.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
index 6f7cadf4..5f97b3b0 100644
--- a/thingsvision/core/extraction/torch.py
+++ b/thingsvision/core/extraction/torch.py
@@ -347,7 +347,6 @@ def __init__(
             raise ValueError(
                 "\nPlease provide either a single module name or a list of module names, but not both.\n"
             )
-        print(module_names, module_name, "INITIAKLAA")
         if module_name is not None:
             module_names = [module_name]
         self.extractor = extractor

From 9c18a3a2c1d5a3636b4aa76f499a61e4e7f39798 Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Wed, 6 Aug 2025 11:55:24 +0200
Subject: [PATCH 18/21] remove print statement

---
 thingsvision/core/extraction/torch.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
index 5f97b3b0..c9e0bcde 100644
--- a/thingsvision/core/extraction/torch.py
+++ b/thingsvision/core/extraction/torch.py
@@ -356,7 +356,6 @@ def __init__(
 
     def __enter__(self) -> PyTorchExtractor:
         """Registering hooks and setting attributes during opening."""
-        print("EXTRACTOR", self.module_names)
         self.extractor._module_and_output_check(self.module_names, self.output_type)
         self.extractor._register_hooks(self.module_names)
         setattr(self.extractor, "module_name", self.module_name)

From a2698ef7e3b52c88a10685a30b24ffc8aaa1741a Mon Sep 17 00:00:00 2001
From: lucaeyring <luca.eyring@googlemail.com>
Date: Wed, 6 Aug 2025 13:34:17 +0200
Subject: [PATCH 19/21] updated setup and requirements

---
 README.md        | 2 --
 requirements.txt | 1 +
 setup.py         | 1 +
 3 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index fc564c8b..819511e1 100644
--- a/README.md
+++ b/README.md
@@ -109,8 +109,6 @@ $ pip install --upgrade thingsvision
 
 The package automatically installs the [Harmonization](https://github.com/serre-lab/harmonization) and [DreamSim](https://github.com/ssundaram21/dreamsim) repositories. See the documentation for available [harmonized models](https://vicco-group.github.io/thingsvision/AvailableModels.html#harmonization) and [DreamSim models](https://vicco-group.github.io/thingsvision/AvailableModels.html#dreamsim) in `thingsvision`.
 
-See the [docs](https://vicco-group.github.io/thingsvision/AvailableModels.html#dreamsim) for which `DreamSim` models are available in `thingsvision`.
-
 #### Google Colab
 Alternatively, you can use Google Colab to play around with `thingsvision` by uploading your image data to Google Drive (via directory mounting).
 You can find the jupyter notebook using `PyTorch` [here](https://colab.research.google.com/github/ViCCo-Group/thingsvision/blob/master/notebooks/pytorch.ipynb) and the `TensorFlow` example [here](https://colab.research.google.com/github/ViCCo-Group/thingsvision/blob/master/notebooks/tensorflow.ipynb).
diff --git a/requirements.txt b/requirements.txt
index 6f5c79b0..46d2c661 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,6 +7,7 @@ numpy<2
 open_clip_torch==3.*
 pandas
 regex
+safetensors<0.6
 scikit-image
 scikit-learn
 scipy
diff --git a/setup.py b/setup.py
index 75f64f94..62c94a67 100644
--- a/setup.py
+++ b/setup.py
@@ -15,6 +15,7 @@
     "open_clip_torch==3.*",
     "pandas",
     "regex",
+    "safetensors<0.6",
     "scikit-image",
     "scikit-learn",
     "scipy",

From 9fd9ad839453d3c4232241f2d8c9e7f7e4987d08 Mon Sep 17 00:00:00 2001
From: Lukas Muttenthaler <LukasMut@users.noreply.github.com>
Date: Wed, 6 Aug 2025 14:07:36 +0200
Subject: [PATCH 20/21] Update README.md

---
 README.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 819511e1..596ea7dc 100644
--- a/README.md
+++ b/README.md
@@ -237,11 +237,12 @@ for batch in my_dataloader:
 
 #### Multi Module Feature Extraction
 
-We've added the functionality to also jointly extract features of multiple `module_names`.
+It is possible to jointly extract features for multiple `module_names` of  a single model.
 
 ##### PyTorch
 
 ```python
+
 module_names = ['visual', ...] # add more module_names here
 
 # your custom dataset and dataloader classes come here (for example, a PyTorch data loader)

From 7d769fce21bbe02e9cc88c621286d40b99edb7dd Mon Sep 17 00:00:00 2001
From: Lukas Muttenthaler <LukasMut@users.noreply.github.com>
Date: Wed, 6 Aug 2025 14:08:50 +0200
Subject: [PATCH 21/21] Update README.md

---
 README.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 596ea7dc..bc13c8f0 100644
--- a/README.md
+++ b/README.md
@@ -16,7 +16,7 @@
         <img src="https://img.shields.io/pypi/dm/thingsvision" alt="downloads">
     </a>
     <a href="https://www.python.org/" rel="nofollow">
-        <img src="https://img.shields.io/badge/python-3.9%20%7C%203.10%20%7C%203.11-blue.svg" alt="Python version" />
+        <img src="https://img.shields.io/badge/python-3.9%20%7C%203.10%20%7C%203.11%20%7C%203.12-blue.svg" alt="Python version" />
     </a>
     <a href="https://github.com/ViCCo-Group/thingsvision/blob/master/LICENSE" rel="nofollow">
         <img src="https://img.shields.io/pypi/l/thingsvision" alt="License" />
@@ -97,7 +97,7 @@ Neural networks come from different sources. With `thingsvision`, you can extrac
 <!-- Setting up your environment -->
 ### :computer: Setting up your environment
 #### Working locally
-First, create a new `conda environment` with Python version 3.10, or 3.11 e.g. by using `conda`:
+First, create a new `conda environment` with Python version 3.10, 3.11, or 3.12 e.g. by using `conda`:
 ```bash
 $ conda create -n thingsvision python=3.10
 $ conda activate thingsvision