diff --git a/README.md b/README.md
index f867421..94a1929 100644
--- a/README.md
+++ b/README.md
@@ -15,11 +15,18 @@ conda create -n tvsd-benchmark
 conda activate tvsd-benchmark
 pip install -r requirements.txt
 ```
+
+For timm model support, ensure timm is installed:
+```bash
+pip install timm
+```
+
 Alternatively, you can use a `venv` environment.
 ```bash
 python -m venv env
 source env/bin/activate
 pip install -r requirements.txt
+pip install timm
 ```
 To obtain the TVSD dataset, run
 ```bash
@@ -54,4 +61,50 @@ Which will generative and evaluate activations for each model.
 
 ## Adding Your Own Model
 
-In the current configuration, each model is specified by a corresponding config file in `configs`. Making a new config for your model is self-explanatory--just follow the outline of the existing ones. You will also have to build out `utils/load_model.py` to accept your added model. In the future, direct integration with `timm` will be provided. 
+In the current configuration, each model is specified by a corresponding config file in `configs`. Making a new config for your model is self-explanatory--just follow the outline of the existing ones. You will also have to build out `utils/load_model.py` to accept your added model.
+
+## Using Timm Models
+
+This repository now has full support for models from the [timm](https://github.com/huggingface/pytorch-image-models) library. To benchmark a timm model:
+
+1. **Create a config file** (or use one of the provided examples in `configs/examples/`):
+
+```yaml
+model-name: resnet50  # Any timm model name
+model-source: timm
+pretrained: true  # Set to false to use random weights
+hook-interval: 8  # Interval for activation extraction
+transform:
+  - name: Resize
+    size: [224, 224]
+  - name: ToTensor
+  - name: Normalize
+    mean: [0.485, 0.456, 0.406]
+    std: [0.229, 0.224, 0.225]
+```
+
+2. **Run benchmarking** as usual:
+
+```bash
+sbatch scripts/generate_activations.sh configs/examples/resnet50.yaml
+sbatch scripts/benchmark.sh configs/examples/resnet50.yaml
+```
+
+### Supported Model Families
+
+The following timm model families are fully supported with specialized activation extraction:
+
+- **ResNet** (`resnet50`, `resnet101`, etc.) - Standard CNN architecture
+- **Vision Transformers** (`vit_base_patch16_224`, `vit_tiny_patch16_224`, etc.) - Handles token/cls shapes
+- **ConvNeXt** (`convnext_base`, `convnext_tiny`, etc.) - Modern CNN with stages
+- **Swin Transformer** (`swin_base_patch4_window7_224`, `swin_tiny_patch4_window7_224`, etc.) - Hierarchical vision transformers
+
+Example configs for these models are provided in `configs/examples/`.
+
+### Note on Model Names
+
+You can list all available timm models using:
+```python
+import timm
+print(timm.list_models(pretrained=True))
+``` 
diff --git a/configs/examples/convnext_base.yaml b/configs/examples/convnext_base.yaml
new file mode 100644
index 0000000..2c7e1d2
--- /dev/null
+++ b/configs/examples/convnext_base.yaml
@@ -0,0 +1,11 @@
+model-name: convnext_base
+model-source: timm
+pretrained: true
+hook-interval: 5
+transform:
+  - name: Resize
+    size: [224, 224]
+  - name: ToTensor
+  - name: Normalize
+    mean: [0.485, 0.456, 0.406]
+    std: [0.229, 0.224, 0.225]
diff --git a/configs/examples/resnet50.yaml b/configs/examples/resnet50.yaml
new file mode 100644
index 0000000..3ae54dd
--- /dev/null
+++ b/configs/examples/resnet50.yaml
@@ -0,0 +1,11 @@
+model-name: resnet50
+model-source: timm
+pretrained: true
+hook-interval: 8
+transform:
+  - name: Resize
+    size: [224, 224]
+  - name: ToTensor
+  - name: Normalize
+    mean: [0.485, 0.456, 0.406]
+    std: [0.229, 0.224, 0.225]
diff --git a/configs/examples/swin_base_patch4_window7_224.yaml b/configs/examples/swin_base_patch4_window7_224.yaml
new file mode 100644
index 0000000..49ad8c4
--- /dev/null
+++ b/configs/examples/swin_base_patch4_window7_224.yaml
@@ -0,0 +1,11 @@
+model-name: swin_base_patch4_window7_224
+model-source: timm
+pretrained: true
+hook-interval: 5
+transform:
+  - name: Resize
+    size: [224, 224]
+  - name: ToTensor
+  - name: Normalize
+    mean: [0.485, 0.456, 0.406]
+    std: [0.229, 0.224, 0.225]
diff --git a/configs/examples/vit_base_patch16_224.yaml b/configs/examples/vit_base_patch16_224.yaml
new file mode 100644
index 0000000..70e4837
--- /dev/null
+++ b/configs/examples/vit_base_patch16_224.yaml
@@ -0,0 +1,11 @@
+model-name: vit_base_patch16_224
+model-source: timm
+pretrained: true
+hook-interval: 5
+transform:
+  - name: Resize
+    size: [224, 224]
+  - name: ToTensor
+  - name: Normalize
+    mean: [0.485, 0.456, 0.406]
+    std: [0.229, 0.224, 0.225]
diff --git a/requirements.txt b/requirements.txt
index 304a8f6..35eb04c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,6 +3,7 @@ scipy
 torch
 torch-pca
 torchvision
+timm
 osfclient
 scikit-learn
 matplotlib
diff --git a/tests/test_timm.py b/tests/test_timm.py
new file mode 100644
index 0000000..92a35d0
--- /dev/null
+++ b/tests/test_timm.py
@@ -0,0 +1,404 @@
+"""Tests for timm model loading and activation extraction."""
+
+import pytest
+import torch
+import numpy as np
+import os
+import tempfile
+import yaml
+from pathlib import Path
+
+from utils.load_model import load_model
+from utils.hooks import Activations
+from utils.timm_helpers import detect_model_family, process_activations
+
+
+class TestTimmModelLoading:
+    """Tests for loading timm models via config."""
+
+    def test_load_timm_resnet(self, tmp_path):
+        """Test loading a timm ResNet model."""
+        config_path = tmp_path / "resnet_config.yaml"
+        config = {
+            "model-name": "resnet50",
+            "model-source": "timm",
+            "pretrained": False,
+            "hook-interval": 5,
+            "transform": [
+                {"name": "Resize", "size": [224, 224]},
+                {"name": "ToTensor"},
+                {"name": "Normalize", "mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]},
+            ],
+        }
+        with open(config_path, "w") as f:
+            yaml.dump(config, f)
+
+        model, model_name, hook_interval = load_model(str(config_path))
+        
+        assert model is not None
+        assert model_name == "resnet50"
+        assert hook_interval == 5
+        assert hasattr(model, "eval")
+
+    def test_load_timm_vit(self, tmp_path):
+        """Test loading a timm Vision Transformer model."""
+        config_path = tmp_path / "vit_config.yaml"
+        config = {
+            "model-name": "vit_tiny_patch16_224",
+            "model-source": "timm",
+            "pretrained": False,
+            "hook-interval": 5,
+            "transform": [
+                {"name": "Resize", "size": [224, 224]},
+                {"name": "ToTensor"},
+                {"name": "Normalize", "mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]},
+            ],
+        }
+        with open(config_path, "w") as f:
+            yaml.dump(config, f)
+
+        model, model_name, hook_interval = load_model(str(config_path))
+        
+        assert model is not None
+        assert model_name == "vit_tiny_patch16_224"
+        assert hook_interval == 5
+
+    def test_load_timm_convnext(self, tmp_path):
+        """Test loading a timm ConvNeXt model."""
+        config_path = tmp_path / "convnext_config.yaml"
+        config = {
+            "model-name": "convnext_tiny",
+            "model-source": "timm",
+            "pretrained": False,
+            "hook-interval": 5,
+            "transform": [
+                {"name": "Resize", "size": [224, 224]},
+                {"name": "ToTensor"},
+                {"name": "Normalize", "mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]},
+            ],
+        }
+        with open(config_path, "w") as f:
+            yaml.dump(config, f)
+
+        model, model_name, hook_interval = load_model(str(config_path))
+        
+        assert model is not None
+        assert model_name == "convnext_tiny"
+        assert hook_interval == 5
+
+    def test_load_timm_swin(self, tmp_path):
+        """Test loading a timm Swin Transformer model."""
+        config_path = tmp_path / "swin_config.yaml"
+        config = {
+            "model-name": "swin_tiny_patch4_window7_224",
+            "model-source": "timm",
+            "pretrained": False,
+            "hook-interval": 5,
+            "transform": [
+                {"name": "Resize", "size": [224, 224]},
+                {"name": "ToTensor"},
+                {"name": "Normalize", "mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]},
+            ],
+        }
+        with open(config_path, "w") as f:
+            yaml.dump(config, f)
+
+        model, model_name, hook_interval = load_model(str(config_path))
+        
+        assert model is not None
+        assert model_name == "swin_tiny_patch4_window7_224"
+        assert hook_interval == 5
+
+
+class TestTimmModelFamilyDetection:
+    """Tests for model family detection."""
+
+    def test_detect_vit(self):
+        """Test detecting Vision Transformer models."""
+        assert detect_model_family("vit_base_patch16_224") == "vit"
+        assert detect_model_family("vit_tiny_patch16_224") == "vit"
+        assert detect_model_family("vision_transformer_small") == "vit"
+
+    def test_detect_swin(self):
+        """Test detecting Swin Transformer models."""
+        assert detect_model_family("swin_base_patch4_window7_224") == "swin"
+        assert detect_model_family("swin_tiny_patch4_window7_224") == "swin"
+
+    def test_detect_convnext(self):
+        """Test detecting ConvNeXt models."""
+        assert detect_model_family("convnext_base") == "convnext"
+        assert detect_model_family("convnext_tiny") == "convnext"
+
+    def test_detect_resnet(self):
+        """Test detecting ResNet models."""
+        assert detect_model_family("resnet50") == "resnet"
+        assert detect_model_family("resnet101") == "resnet"
+
+    def test_detect_default(self):
+        """Test default detection for unknown models."""
+        assert detect_model_family("some_other_model") == "default"
+
+
+class TestTimmActivationExtraction:
+    """Tests for activation extraction from timm models."""
+
+    @pytest.fixture
+    def dummy_batch(self):
+        """Create a dummy batch of images."""
+        return torch.randn(2, 3, 224, 224)
+
+    @pytest.fixture
+    def output_dir(self, tmp_path):
+        """Create a temporary output directory."""
+        return str(tmp_path)
+
+    def test_resnet_activations(self, dummy_batch, output_dir):
+        """Test activation extraction from ResNet (easy baseline)."""
+        import timm
+        model = timm.create_model("resnet50", pretrained=False)
+        model.eval()
+
+        activations = Activations(
+            output_dir=output_dir,
+            model_name="resnet50",
+            dataset_name="test",
+            pca_components=None,
+        )
+
+        # Register hooks on a few layers
+        layer_names = ["layer1.0", "layer2.0", "layer3.0"]
+        activations.register(model, layer_names)
+        
+        # Forward pass
+        activations.set_batch(0)
+        activations.set_training_mode(False)
+        with torch.no_grad():
+            _ = model(dummy_batch)
+        activations.finalize_batch_inference()
+
+        # Check activations were captured
+        assert len(activations.activations) == len(layer_names)
+        for layer_name in layer_names:
+            assert layer_name in activations.activations
+            assert 0 in activations.activations[layer_name]
+            # Check shape: should be [batch_size, features]
+            act_tensor = activations.activations[layer_name][0][0]
+            assert act_tensor.shape[0] == dummy_batch.shape[0]
+            assert len(act_tensor.shape) == 2
+
+    def test_vit_activations(self, dummy_batch, output_dir):
+        """Test activation extraction from Vision Transformer (token/cls shapes)."""
+        import timm
+        model = timm.create_model("vit_tiny_patch16_224", pretrained=False)
+        model.eval()
+
+        activations = Activations(
+            output_dir=output_dir,
+            model_name="vit_tiny_patch16_224",
+            dataset_name="test",
+            pca_components=None,
+        )
+
+        # Dynamically find transformer block layers
+        block_layers = [name for name, _ in model.named_modules() if name.startswith("blocks.")]
+        # Select a few blocks (but ensure they exist)
+        if len(block_layers) >= 3:
+            layer_names = [block_layers[0], block_layers[len(block_layers)//2], block_layers[-1]]
+        else:
+            layer_names = block_layers[:min(3, len(block_layers))]
+        
+        if not layer_names:
+            pytest.skip("No transformer blocks found in model")
+        
+        activations.register(model, layer_names)
+        
+        # Forward pass
+        activations.set_batch(0)
+        activations.set_training_mode(False)
+        with torch.no_grad():
+            _ = model(dummy_batch)
+        activations.finalize_batch_inference()
+
+        # Check activations were captured
+        assert len(activations.activations) > 0
+        for layer_name in activations.activations.keys():
+            assert 0 in activations.activations[layer_name]
+            # Check shape: should be [batch_size, features] after processing
+            act_tensor = activations.activations[layer_name][0][0]
+            assert act_tensor.shape[0] == dummy_batch.shape[0]
+            assert len(act_tensor.shape) == 2
+
+    def test_convnext_activations(self, dummy_batch, output_dir):
+        """Test activation extraction from ConvNeXt (stages)."""
+        import timm
+        model = timm.create_model("convnext_tiny", pretrained=False)
+        model.eval()
+
+        activations = Activations(
+            output_dir=output_dir,
+            model_name="convnext_tiny",
+            dataset_name="test",
+            pca_components=None,
+        )
+
+        # Register hooks on stages
+        layer_names = ["stages.0", "stages.1", "stages.2"]
+        activations.register(model, layer_names)
+        
+        # Forward pass
+        activations.set_batch(0)
+        activations.set_training_mode(False)
+        with torch.no_grad():
+            _ = model(dummy_batch)
+        activations.finalize_batch_inference()
+
+        # Check activations were captured
+        assert len(activations.activations) > 0
+        for layer_name in activations.activations.keys():
+            assert 0 in activations.activations[layer_name]
+            # Check shape: should be [batch_size, features]
+            act_tensor = activations.activations[layer_name][0][0]
+            assert act_tensor.shape[0] == dummy_batch.shape[0]
+            assert len(act_tensor.shape) == 2
+
+    def test_swin_activations(self, dummy_batch, output_dir):
+        """Test activation extraction from Swin Transformer (hierarchical tokens)."""
+        import timm
+        model = timm.create_model("swin_tiny_patch4_window7_224", pretrained=False)
+        model.eval()
+
+        activations = Activations(
+            output_dir=output_dir,
+            model_name="swin_tiny_patch4_window7_224",
+            dataset_name="test",
+            pca_components=None,
+        )
+
+        # Dynamically find layer modules
+        swin_layers = [name for name, _ in model.named_modules() if name.startswith("layers.")]
+        # Select available layers (up to 3)
+        layer_names = swin_layers[:min(3, len(swin_layers))]
+        
+        if not layer_names:
+            pytest.skip("No Swin layers found in model")
+        
+        activations.register(model, layer_names)
+        
+        # Forward pass
+        activations.set_batch(0)
+        activations.set_training_mode(False)
+        with torch.no_grad():
+            _ = model(dummy_batch)
+        activations.finalize_batch_inference()
+
+        # Check activations were captured
+        assert len(activations.activations) > 0
+        for layer_name in activations.activations.keys():
+            assert 0 in activations.activations[layer_name]
+            # Check shape: should be [batch_size, features]
+            act_tensor = activations.activations[layer_name][0][0]
+            assert act_tensor.shape[0] == dummy_batch.shape[0]
+            assert len(act_tensor.shape) == 2
+
+    def test_activation_shape_processing(self):
+        """Test that activation processing handles different output shapes correctly."""
+        batch_size = 4
+        
+        # Test ViT-style output (batch, num_tokens, embed_dim)
+        vit_output = torch.randn(batch_size, 197, 768)  # 196 patches + 1 cls token
+        processed = process_activations(vit_output, "vit")
+        assert processed.shape == (batch_size, 197 * 768)
+        
+        # Test ConvNeXt-style output (batch, channels, H, W)
+        convnext_output = torch.randn(batch_size, 96, 56, 56)
+        processed = process_activations(convnext_output, "convnext")
+        assert processed.shape == (batch_size, 96 * 56 * 56)
+        
+        # Test Swin-style output (batch, H*W, embed_dim)
+        swin_output = torch.randn(batch_size, 3136, 96)  # 56*56 patches
+        processed = process_activations(swin_output, "swin")
+        assert processed.shape == (batch_size, 3136 * 96)
+
+    def test_no_crashes_on_forward_pass(self, dummy_batch, output_dir):
+        """Test that models don't crash during forward pass with hooks."""
+        import timm
+        
+        models_to_test = [
+            "resnet50",
+            "vit_tiny_patch16_224",
+            "convnext_tiny",
+            "swin_tiny_patch4_window7_224",
+        ]
+        
+        for model_name in models_to_test:
+            model = timm.create_model(model_name, pretrained=False)
+            model.eval()
+            
+            activations = Activations(
+                output_dir=output_dir,
+                model_name=model_name,
+                dataset_name="test",
+                pca_components=None,
+            )
+            
+            # Get some layer names
+            layer_names = [name for name, _ in model.named_modules()]
+            # Take every 50th layer to avoid too many hooks
+            sample_layers = layer_names[::50][:3]
+            
+            if sample_layers:
+                activations.register(model, sample_layers)
+                activations.set_batch(0)
+                activations.set_training_mode(False)
+                
+                # This should not crash
+                with torch.no_grad():
+                    _ = model(dummy_batch)
+                activations.finalize_batch_inference()
+
+
+class TestTimmActivationSaving:
+    """Tests for saving timm model activations."""
+
+    @pytest.fixture
+    def dummy_batch(self):
+        """Create a dummy batch of images."""
+        return torch.randn(2, 3, 224, 224)
+
+    @pytest.fixture
+    def output_dir(self, tmp_path):
+        """Create a temporary output directory."""
+        return str(tmp_path)
+
+    def test_save_resnet_activations(self, dummy_batch, output_dir):
+        """Test saving ResNet activations to disk."""
+        import timm
+        model = timm.create_model("resnet50", pretrained=False)
+        model.eval()
+
+        activations = Activations(
+            output_dir=output_dir,
+            model_name="resnet50",
+            dataset_name="test",
+            pca_components=None,
+        )
+
+        layer_names = ["layer1.0"]
+        activations.register(model, layer_names)
+        
+        activations.set_batch(0)
+        activations.set_training_mode(False)
+        with torch.no_grad():
+            _ = model(dummy_batch)
+        activations.finalize_batch_inference()
+        
+        # Save and check file exists
+        activations.save()
+        
+        expected_path = Path(output_dir) / "activations" / "test" / "resnet50" / "layer1.0" / "activations.pt"
+        assert expected_path.exists()
+        
+        # Load and verify
+        loaded = torch.load(expected_path)
+        assert loaded.shape[0] == dummy_batch.shape[0]
+        assert len(loaded.shape) == 2
diff --git a/utils/hooks.py b/utils/hooks.py
index ba3e131..acf73e1 100644
--- a/utils/hooks.py
+++ b/utils/hooks.py
@@ -5,6 +5,8 @@
 from sklearn.decomposition import IncrementalPCA
 from typing import List, Dict, Optional
 
+from utils.timm_helpers import detect_model_family, process_activations
+
 
 class Activations:
     def __init__(
@@ -25,6 +27,8 @@ def __init__(
         self._training_mode = False
         self.pca_components = pca_components
         self._training_activations = {}  # {layer_name: [activations]}
+        # Detect model family for proper activation processing
+        self.model_family = detect_model_family(model_name)
 
     def set_batch(self, batch):
         self._current_batch = batch
@@ -35,8 +39,8 @@ def set_training_mode(self, training: bool):
     def _get_hook(self, layer_name, debug=False):
         def hook_fn(module, input, output):
             if self._active:
-                if isinstance(output, list) or isinstance(output, tuple):
-                    output = torch.stack(output, dim=0)
+                # Use timm_helpers to properly process activations
+                output = process_activations(output, self.model_family)
 
                 if self._training_mode:
                     self._handle_training_mode(layer_name, output)
diff --git a/utils/load_model.py b/utils/load_model.py
index a66cc94..1c53844 100644
--- a/utils/load_model.py
+++ b/utils/load_model.py
@@ -1,14 +1,21 @@
 import sys
 import yaml
 
-# import timm
+import timm
 import torch
 from torchvision import transforms
 
-sys.path.append(
-    "/users/jamullik/pytorch-image-models/"
-)  # hacky for now, need to fix this later
-from timm.models.RESMAX import chresmax_v3
+# Only import RESMAX if available (for hmax models)
+# Note: This path is user-specific. Set the RESMAX_PATH environment variable
+# or modify this path for your environment if using HMAX models.
+try:
+    import os
+    resmax_path = os.environ.get('RESMAX_PATH', '/users/jamullik/pytorch-image-models/')
+    sys.path.append(resmax_path)
+    from timm.models.RESMAX import chresmax_v3
+    HAS_RESMAX = True
+except ImportError:
+    HAS_RESMAX = False
 
 
 def load_model(config_path: str):
@@ -64,12 +71,18 @@ def load_torchvision_model(config: dict):
 
 
 def load_timm_model(config: dict):
-    model = timm.create_model(config["model-name"], pretrained=True)
+    pretrained = config.get("pretrained", True)
+    model = timm.create_model(config["model-name"], pretrained=pretrained)
     model_name = config["model-name"]
     return model, model_name
 
 
 def load_hmax_model(config: dict):
+    if not HAS_RESMAX:
+        raise ImportError(
+            "RESMAX module is not available. Please check your installation."
+        )
+    
     if config["model-type"] == "chresmax_v3":
         checkpoint = torch.load(
             config["hmax-info"]["ckpt_path"],
diff --git a/utils/timm_helpers.py b/utils/timm_helpers.py
new file mode 100644
index 0000000..f66a954
--- /dev/null
+++ b/utils/timm_helpers.py
@@ -0,0 +1,111 @@
+"""Helper utilities for handling timm model activations."""
+
+import torch
+
+
+def detect_model_family(model_name: str) -> str:
+    """
+    Detect the model family from the model name.
+    
+    Args:
+        model_name: Name of the model
+        
+    Returns:
+        Model family name: 'vit', 'swin', 'convnext', 'resnet', or 'default'
+    """
+    model_name_lower = model_name.lower()
+    
+    if 'vit' in model_name_lower or 'vision_transformer' in model_name_lower:
+        return 'vit'
+    elif 'swin' in model_name_lower:
+        return 'swin'
+    elif 'convnext' in model_name_lower:
+        return 'convnext'
+    elif 'resnet' in model_name_lower:
+        return 'resnet'
+    else:
+        return 'default'
+
+
+def process_vit_activations(output: torch.Tensor) -> torch.Tensor:
+    """
+    Process Vision Transformer activations.
+    
+    ViT models have outputs with shape (batch, num_tokens, embed_dim)
+    where num_tokens = num_patches + 1 (cls token).
+    We flatten the token and embedding dimensions.
+    
+    Args:
+        output: Activation tensor from a ViT layer
+        
+    Returns:
+        Processed tensor with shape (batch, num_tokens * embed_dim)
+    """
+    if len(output.shape) == 3:  # (batch, num_tokens, embed_dim)
+        # Flatten token and embedding dimensions
+        batch_size = output.shape[0]
+        return output.reshape(batch_size, -1)
+    return output
+
+
+def process_swin_activations(output: torch.Tensor) -> torch.Tensor:
+    """
+    Process Swin Transformer activations.
+    
+    Swin models have hierarchical outputs with varying shapes depending on stage.
+    Early stages: (batch, H, W, embed_dim)
+    Later stages: (batch, num_patches, embed_dim)
+    
+    Args:
+        output: Activation tensor from a Swin layer
+        
+    Returns:
+        Processed tensor with shape (batch, features)
+    """
+    batch_size = output.shape[0]
+    return output.reshape(batch_size, -1)
+
+
+def process_convnext_activations(output: torch.Tensor) -> torch.Tensor:
+    """
+    Process ConvNeXt activations.
+    
+    ConvNeXt models have standard conv outputs with shape (batch, channels, H, W).
+    
+    Args:
+        output: Activation tensor from a ConvNeXt layer
+        
+    Returns:
+        Processed tensor with shape (batch, channels * H * W)
+    """
+    batch_size = output.shape[0]
+    return output.reshape(batch_size, -1)
+
+
+def process_activations(output: torch.Tensor, model_family: str) -> torch.Tensor:
+    """
+    Process activations based on model family.
+    
+    Args:
+        output: Raw activation tensor
+        model_family: Type of model ('vit', 'swin', 'convnext', 'resnet', 'default')
+        
+    Returns:
+        Processed activation tensor with shape (batch, features)
+    """
+    # Handle tuple/list outputs (e.g., from some model layers)
+    if isinstance(output, (list, tuple)):
+        # Take the first element if multiple outputs are provided
+        output = output[0]
+    
+    # Apply family-specific processing
+    if model_family == 'vit':
+        return process_vit_activations(output)
+    elif model_family == 'swin':
+        return process_swin_activations(output)
+    elif model_family == 'convnext':
+        return process_convnext_activations(output)
+    else:
+        # Default processing: flatten all dimensions except batch
+        batch_size = output.shape[0]
+        return output.reshape(batch_size, -1)