Bug fix: type issue with models leading to gibberish outputs

app.py

@@ -16,7 +16,8 @@ import dash
 from dash import html, dcc, Input, Output, State, callback, no_update, ALL, MATCH
 import json
 import torch
-from utils import (load_model_and_get_patterns, execute_forward_pass, extract_layer_data,
+from utils import (load_model_for_inference, load_model_and_get_patterns,
+                   execute_forward_pass, extract_layer_data,
                    perform_beam_search, execute_forward_pass_with_multi_layer_head_ablation)
 from utils.head_detection import get_active_head_summary
 from utils.model_config import get_auto_selections
@@ -374,10 +375,9 @@ def run_generation(n_clicks, model_name, prompt, max_new_tokens, beam_width, pat
         return no_update, no_update, no_update, no_update, no_update, no_update, no_update, no_update
 
     try:
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        model = AutoModelForCausalLM.from_pretrained(model_name, attn_implementation='eager')
+        from transformers import AutoTokenizer
+        model = load_model_for_inference(model_name)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model.eval()
 
         # Always run beam search (even with max_new_tokens=1)
         results = perform_beam_search(model, tokenizer, prompt, beam_width, max_new_tokens)
@@ -555,16 +555,15 @@ def store_selected_beam(n_clicks_list, results_data, existing_activation_data, o
     new_activation_data = no_update
     if existing_activation_data:
         try:
-            from transformers import AutoModelForCausalLM, AutoTokenizer
+            from transformers import AutoTokenizer
             model_name = existing_activation_data['model']
             config = {
                 'attention_modules': existing_activation_data['attention_modules'],
                 'block_modules': existing_activation_data['block_modules'],
                 'norm_parameters': existing_activation_data.get('norm_parameters', [])
             }
-            model = AutoModelForCausalLM.from_pretrained(model_name, attn_implementation='eager')
+            model = load_model_for_inference(model_name)
             tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model.eval()
             # Pass original_prompt so per-position top-5 data is computed for scrubber
             orig_prompt = original_prompt_data.get('prompt', '') if original_prompt_data else ''
             new_activation_data = execute_forward_pass(
@@ -613,10 +612,10 @@ def update_pipeline_content(activation_data, model_name):
         return tuple(empty_outputs)
     
     try:
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        model = AutoModelForCausalLM.from_pretrained(model_name, attn_implementation='eager')
+        from transformers import AutoTokenizer
+        model = load_model_for_inference(model_name)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        
+
         # Use pre-decoded tokens if available, otherwise decode from input_ids
         input_ids = activation_data.get('input_ids', [[]])[0]
         tokens = activation_data.get('tokens') or [tokenizer.decode([tid]) for tid in input_ids]
@@ -922,11 +921,10 @@ def run_ablation_experiment(n_clicks, selected_heads, activation_data, model_nam
         return no_update, no_update, no_update
     
     try:
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        model = AutoModelForCausalLM.from_pretrained(model_name, attn_implementation='eager')
+        from transformers import AutoTokenizer
+        model = load_model_for_inference(model_name)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model.eval()
-        
+
         sequence_text = prompt
         
         config = {
@@ -1087,7 +1085,9 @@ def update_attribution_target_options(activation_data):
     options = []
     for t in global_top5:
         if isinstance(t, dict):
-            options.append({'label': f"{t['token']} ({t['probability']:.1%})", 'value': t['token']})
+            prob = t.get('probability')
+            prob_str = f" ({prob:.1%})" if prob is not None else ""
+            options.append({'label': f"{t['token']}{prob_str}", 'value': t['token']})
         else:
             options.append({'label': t[0], 'value': t[0]})
     return options
@@ -1108,11 +1108,10 @@ def run_attribution_experiment(n_clicks, method, target_token, activation_data,
         return no_update
     
     try:
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        model = AutoModelForCausalLM.from_pretrained(model_name, attn_implementation='eager')
+        from transformers import AutoTokenizer
+        model = load_model_for_inference(model_name)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model.eval()
-        
+
         sequence_text = activation_data.get('prompt', prompt)
         
         # Get target token ID if specified

components/pipeline.py

@@ -1127,7 +1127,7 @@ def create_output_content(top_tokens=None, predicted_token=None, predicted_prob=
         fig = go.Figure(go.Bar(
             x=probs, y=tokens, orientation='h',
             marker_color=['#00f2fe' if i == 0 else '#4facfe' for i in range(len(tokens))],
-            text=[f"{p:.1%}" for p in probs], textposition='outside',
+            text=[f"{p:.1%}" if p is not None else "?" for p in probs], textposition='outside',
             hovertemplate='%{y} (%{x:.1%})<extra></extra>'
         ))
         fig.update_layout(

debug_logs.md

@@ -0,0 +1,163 @@
+-------------------------------------------------------
+Qwen2ForCausalLM LOAD REPORT from: Qwen/Qwen2.5-0.5B
+Key            | Status  | 
+---------------+---------+-
+lm_head.weight | MISSING | 
+
+Notes:
+- MISSING	:those params were newly initialized because missing from the checkpoint. Consider training on your downstream task.
+-------------------------------------------------------
+Traceback (most recent call last):
+  File "/app/app.py", line 383, in run_generation
+    results = perform_beam_search(model, tokenizer, prompt, beam_width, max_new_tokens)
+  File "/app/utils/beam_search.py", line 142, in perform_beam_search
+    outputs = model(seq)
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
+    return forward_call(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/transformers/utils/generic.py", line 843, in wrapper
+    output = func(self, *args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 476, in forward
+    outputs: BaseModelOutputWithPast = self.model(
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
+    return forward_call(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/transformers/utils/generic.py", line 917, in wrapper
+    output = func(self, *args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/transformers/utils/output_capturing.py", line 253, in wrapper
+    outputs = func(self, *args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 411, in forward
+    hidden_states = decoder_layer(
+  File "/usr/local/lib/python3.10/site-packages/transformers/modeling_layers.py", line 93, in __call__
+    return super().__call__(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
+    return forward_call(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 298, in forward
+    hidden_states, _ = self.self_attn(
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
+    return forward_call(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 218, in forward
+    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
+    return forward_call(*args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 134, in forward
+    return F.linear(input, self.weight, self.bias)
+RuntimeError: mat1 and mat2 must have the same dtype, but got Float and BFloat16
+-------------------------------------------------------
+Qwen2ForCausalLM LOAD REPORT from: Qwen/Qwen2.5-0.5B
+Key            | Status  | 
+---------------+---------+-
+lm_head.weight | MISSING | 
+
+Notes:
+- MISSING	:those params were newly initialized because missing from the checkpoint. Consider training on your downstream task.
+10.16.43.195 - - [19/Mar/2026 19:47:47] "POST /_dash-update-component HTTP/1.1" 200 -
+10.16.43.195 - - [19/Mar/2026 19:47:55] "POST /_dash-update-component HTTP/1.1" 200 -
+10.16.43.195 - - [19/Mar/2026 19:47:55] "POST /_dash-update-component HTTP/1.1" 200 -
+Executing forward pass with prompt: 'Draw ascii art for a cat'
+Captured 48 module outputs using PyVene
+Loading weights:   0%|          | 0/290 [00:00<?, ?it/s]10.16.31.44 - - [19/Mar/2026 19:47:55] "POST /_dash-update-component HTTP/1.1" 200 -
+[2026-03-19 19:47:55,972] ERROR in app: Exception on /_dash-update-component [POST]
+Traceback (most recent call last):
+  File "/usr/local/lib/python3.10/site-packages/flask/app.py", line 1511, in wsgi_app
+    response = self.full_dispatch_request()
+  File "/usr/local/lib/python3.10/site-packages/flask/app.py", line 919, in full_dispatch_request
+    rv = self.handle_user_exception(e)
+  File "/usr/local/lib/python3.10/site-packages/flask/app.py", line 917, in full_dispatch_request
+    rv = self.dispatch_request()
+  File "/usr/local/lib/python3.10/site-packages/flask/app.py", line 902, in dispatch_request
+    return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args)  # type: ignore[no-any-return]
+  File "/usr/local/lib/python3.10/site-packages/dash/_get_app.py", line 17, in wrap
+    return ctx.run(func, self, *args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/dash/dash.py", line 1600, in dispatch
+    response_data = ctx.run(partial_func)
+  File "/usr/local/lib/python3.10/site-packages/dash/_callback.py", line 720, in add_context
+    raise err
+  File "/usr/local/lib/python3.10/site-packages/dash/_callback.py", line 711, in add_context
+    output_value = _invoke_callback(func, *func_args, **func_kwargs)  # type: ignore[reportArgumentType]
+  File "/usr/local/lib/python3.10/site-packages/dash/_callback.py", line 58, in _invoke_callback
+    return func(*args, **kwargs)  # %% callback invoked %%
+  File "/app/app.py", line 1090, in update_attribution_target_options
+    options.append({'label': f"{t['token']} ({t['probability']:.1%})", 'value': t['token']})
+TypeError: unsupported format string passed to NoneType.__format__
+-------------------------------------------------------
+DEBUG extract_layer_data: Found 24 attention modules
+Loading model: gpt2-medium
+Loading weights:   0%|          | 0/292 [00:00<?, ?it/s]
+Loading weights:  51%|█████     | 149/292 [00:00<00:00, 1299.04it/s]
+Loading weights: 100%|██████████| 292/292 [00:00<00:00, 1450.02it/s]
+GPT2LMHeadModel LOAD REPORT from: gpt2-medium
+Key            | Status  | 
+---------------+---------+-
+lm_head.weight | MISSING | 
+
+Notes:
+- MISSING	:those params were newly initialized because missing from the checkpoint. Consider training on your downstream task.
+-------------------------------------------------------
+10.16.43.195 - - [19/Mar/2026 20:11:03] "POST /_dash-update-component HTTP/1.1" 200 -
+[2026-03-19 20:11:03,238] ERROR in app: Exception on /_dash-update-component [POST]
+Traceback (most recent call last):
+  File "/usr/local/lib/python3.10/site-packages/flask/app.py", line 1511, in wsgi_app
+    response = self.full_dispatch_request()
+  File "/usr/local/lib/python3.10/site-packages/flask/app.py", line 919, in full_dispatch_request
+    rv = self.handle_user_exception(e)
+  File "/usr/local/lib/python3.10/site-packages/flask/app.py", line 917, in full_dispatch_request
+    rv = self.dispatch_request()
+  File "/usr/local/lib/python3.10/site-packages/flask/app.py", line 902, in dispatch_request
+    return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args)  # type: ignore[no-any-return]
+  File "/usr/local/lib/python3.10/site-packages/dash/_get_app.py", line 17, in wrap
+    return ctx.run(func, self, *args, **kwargs)
+  File "/usr/local/lib/python3.10/site-packages/dash/dash.py", line 1600, in dispatch
+    response_data = ctx.run(partial_func)
+  File "/usr/local/lib/python3.10/site-packages/dash/_callback.py", line 720, in add_context
+    raise err
+  File "/usr/local/lib/python3.10/site-packages/dash/_callback.py", line 711, in add_context
+    output_value = _invoke_callback(func, *func_args, **func_kwargs)  # type: ignore[reportArgumentType]
+  File "/usr/local/lib/python3.10/site-packages/dash/_callback.py", line 58, in _invoke_callback
+    return func(*args, **kwargs)  # %% callback invoked %%
+  File "/app/app.py", line 1090, in update_attribution_target_options
+    options.append({'label': f"{t['token']} ({t['probability']:.1%})", 'value': t['token']})
+TypeError: unsupported format string passed to NoneType.__format__
+10.16.43.195 - - [19/Mar/2026 20:11:03] "POST /_dash-update-component HTTP/1.1" 500 -
+10.16.43.195 - - [19/Mar/2026 20:11:03] "POST /_dash-update-component HTTP/1.1" 200 -
+Traceback (most recent call last):
+  File "/app/utils/model_patterns.py", line 1337, in generate_bertviz_html
+    attention_weights = torch.tensor(attention_output[1])  # [batch, heads, seq, seq]
+RuntimeError: Could not infer dtype of NoneType
+Traceback (most recent call last):
+  File "/app/app.py", line 691, in update_pipeline_content
+    outputs.append(create_output_content(
+  File "/app/components/pipeline.py", line 1130, in create_output_content
+    text=[f"{p:.1%}" for p in probs], textposition='outside',
+  File "/app/components/pipeline.py", line 1130, in <listcomp>
+    text=[f"{p:.1%}" for p in probs], textposition='outside',
+TypeError: unsupported format string passed to NoneType.__format__
+-------------------------------------------------------
+10.16.43.195 - - [19/Mar/2026 20:26:26] "POST /_dash-update-component HTTP/1.1" 200 -
+DEBUG extract_layer_data: Found 24 attention modules
+Warning: Could not compute logit lens for gpt_neox.layers.0: mixed dtype (CPU): expect parameter to have scalar type of Float
+Warning: Could not compute token probabilities for gpt_neox.layers.0: mixed dtype (CPU): expect parameter to have scalar type of Float
+Warning: Could not compute logit lens for gpt_neox.layers.1: mixed dtype (CPU): expect parameter to have scalar type of Float
+Warning: Could not compute token probabilities for gpt_neox.layers.1: mixed dtype (CPU): expect parameter to have scalar type of Float
+Warning: Could not compute logit lens for gpt_neox.layers.2: mixed dtype (CPU): expect parameter to have scalar type of Float
+Warning: Could not compute token probabilities for gpt_neox.layers.2: mixed dtype (CPU): expect parameter to have scalar type of Float
+-------------------------------------------------------
+Warning: Could not compute logit lens for gpt_neox.layers.13: mixed dtype (CPU): expect parameter to have scalar type of Float
+Warning: Could not compute token probabilities for gpt_neox.layers.13: mixed dtype (CPU): expect parameter to have scalar type of Float
+Warning: Could not compute logit lens for gpt_neox.layers.14: Could not infer dtype of NoneType
+Warning: Could not compute token probabilities for gpt_neox.layers.14: Could not infer dtype of NoneType
+Warning: Could not compute logit lens for gpt_neox.layers.15: Could not infer dtype of NoneType
+Warning: Could not compute token probabilities for gpt_neox.layers.15: Could not infer dtype of NoneType
+Warning: Could not compute logit lens for gpt_neox.layers.16: Could not infer dtype of NoneType
+Warning: Could not compute token probabilities for gpt_neox.layers.16: Could not infer dtype of NoneType
+Warning: Could not compute logit lens for gpt_neox.layers.17: Could not infer dtype of NoneType
+Warning: Could not compute token probabilities for gpt_neox.layers.17: Could not infer dtype of NoneType
+-------------------------------------------------------

tests/test_model_patterns.py

@@ -10,8 +10,8 @@ Tests pure logic functions that don't require model loading:
 import pytest
 import torch
 import numpy as np
-from utils.model_patterns import merge_token_probabilities, safe_to_serializable
-from utils import execute_forward_pass_with_multi_layer_head_ablation
+from utils.model_patterns import merge_token_probabilities, safe_to_serializable, _prepare_hidden_state
+from utils import execute_forward_pass_with_multi_layer_head_ablation, load_model_for_inference
 
 
 class TestMergeTokenProbabilities:
@@ -478,3 +478,43 @@ class TestFullSequenceAttentionData:
             attn = data['attention_outputs'][module]['output'][1]
             assert len(attn[0][0]) == 8
             assert len(attn[0][0][0]) == 8
+
+
+class TestPrepareHiddenState:
+    """Tests for _prepare_hidden_state helper."""
+
+    def test_raises_on_none(self):
+        """_prepare_hidden_state(None) should raise ValueError."""
+        with pytest.raises(ValueError, match="Layer output is None"):
+            _prepare_hidden_state(None)
+
+    def test_unwraps_tuple_with_none_second(self):
+        """Tuple where second element is None should unwrap first element."""
+        result = _prepare_hidden_state(([[1.0, 2.0]], None))
+        assert isinstance(result, torch.Tensor)
+        assert result.shape[-1] == 2
+
+    def test_converts_list(self):
+        """Plain list should be converted to torch.Tensor."""
+        result = _prepare_hidden_state([[[1.0, 2.0]]])
+        assert isinstance(result, torch.Tensor)
+
+
+class TestSafeToSerializableTupleWithNone:
+    """Test that safe_to_serializable handles tuples containing None."""
+
+    def test_tuple_with_tensor_and_none(self):
+        """Tuple of (tensor, None) should become [list, None]."""
+        tensor = torch.tensor([1.0, 2.0])
+        result = safe_to_serializable((tensor, None))
+        assert isinstance(result, list)
+        assert result[0] == [1.0, 2.0]
+        assert result[1] is None
+
+
+class TestLoadModelForInference:
+    """Tests for load_model_for_inference helper."""
+
+    def test_function_is_importable(self):
+        """load_model_for_inference should be importable from utils."""
+        assert callable(load_model_for_inference)

utils/__init__.py

@@ -1,11 +1,12 @@
-from .model_patterns import (load_model_and_get_patterns, execute_forward_pass, 
-                             logit_lens_transformation, extract_layer_data, 
+from .model_patterns import (load_model_for_inference, load_model_and_get_patterns,
+                             execute_forward_pass,
+                             logit_lens_transformation, extract_layer_data,
                              generate_bertviz_html,
                              execute_forward_pass_with_head_ablation,
                              execute_forward_pass_with_multi_layer_head_ablation,
-                             merge_token_probabilities, 
+                             merge_token_probabilities,
                              compute_global_top5_tokens, compute_per_position_top5,
-                             detect_significant_probability_increases, 
+                             detect_significant_probability_increases,
                              evaluate_sequence_ablation, generate_bertviz_model_view_html)
 from .model_config import get_model_family, get_family_config, get_auto_selections, MODEL_TO_FAMILY, MODEL_FAMILIES
 from .head_detection import load_head_categories, verify_head_activation, get_active_head_summary
@@ -16,7 +17,8 @@ from .token_attribution import compute_integrated_gradients, compute_simple_grad
 
 __all__ = [
     # Model patterns
-    'load_model_and_get_patterns', 
+    'load_model_for_inference',
+    'load_model_and_get_patterns',
     'execute_forward_pass',
     'execute_forward_pass_with_head_ablation',
     'execute_forward_pass_with_multi_layer_head_ablation',

utils/model_patterns.py

@@ -7,6 +7,26 @@ from typing import Dict, List, Tuple, Any, Optional
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
+def load_model_for_inference(model_name: str):
+    """Load model with float32 dtype for CPU stability and verify weight tying."""
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        attn_implementation='eager',
+        torch_dtype=torch.float32
+    )
+    model.eval()
+
+    # Verify lm_head is properly tied to embeddings (not randomly initialized)
+    embed = model.get_input_embeddings()
+    lm_head = model.get_output_embeddings()
+    if embed is not None and lm_head is not None:
+        if embed.weight.data_ptr() != lm_head.weight.data_ptr():
+            print(f"Warning: {model_name} lm_head not tied to embeddings, re-tying...")
+            model.tie_weights()
+
+    return model
+
+
 def extract_patterns(model, use_modules=True) -> Dict[str, List[str]]:
     """Extract patterns from model modules or parameters."""
     items = model.named_modules() if use_modules else model.named_parameters()
@@ -36,9 +56,8 @@ def load_model_and_get_patterns(model_name: str) -> Tuple[Dict[str, List[str]],
     print(f"Loading model: {model_name}")
     
     # Load model and tokenizer
-    model = AutoModelForCausalLM.from_pretrained(model_name, attn_implementation='eager')
+    model = load_model_for_inference(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model.eval()
     
     # Extract patterns
     module_patterns = extract_patterns(model, use_modules=True)
@@ -919,10 +938,13 @@ def evaluate_sequence_ablation(model, tokenizer, sequence_text: str, config: Dic
 
 def _prepare_hidden_state(layer_output: Any) -> torch.Tensor:
     """Helper to convert layer output to tensor, handling tuple outputs."""
+    if layer_output is None:
+        raise ValueError("Layer output is None")
+
     # Handle PyVene captured tuple outputs where 2nd element is None (e.g. use_cache=False)
     if isinstance(layer_output, (list, tuple)) and len(layer_output) > 1 and layer_output[1] is None:
         layer_output = layer_output[0]
-        
+
     hidden = torch.tensor(layer_output) if not isinstance(layer_output, torch.Tensor) else layer_output
     if hidden.dim() == 4:
         hidden = hidden.squeeze(0)
@@ -954,7 +976,9 @@ def logit_lens_transformation(layer_output: Any, norm_data: List[Any], model, to
     with torch.no_grad():
         # Convert to tensor and ensure proper shape [batch, seq_len, hidden_dim]
         hidden = _prepare_hidden_state(layer_output)
-        
+        # Serialized intermediates may be float64; cast to model dtype
+        hidden = hidden.to(dtype=next(model.parameters()).dtype)
+
         # Step 1: Apply final layer normalization (critical for intermediate layers)
         final_norm = get_norm_layer_from_parameter(model, norm_parameter)
         if final_norm is not None:
@@ -1271,7 +1295,10 @@ def generate_bertviz_model_view_html(activation_data: Dict[str, Any]) -> str:
                 attention_output = attention_outputs[module_name]['output']
                 if isinstance(attention_output, list) and len(attention_output) >= 2:
                     # Get attention weights (element 1 of the output tuple)
-                    attention_weights = torch.tensor(attention_output[1])  # [batch, heads, seq, seq]
+                    raw_weights = attention_output[1]
+                    if raw_weights is None:
+                        continue  # Skip layers with missing attention data
+                    attention_weights = torch.tensor(raw_weights)  # [batch, heads, seq, seq]
                     layer_attention_pairs.append((layer_num, attention_weights))
         
         if not layer_attention_pairs:
@@ -1334,7 +1361,10 @@ def generate_bertviz_html(activation_data: Dict[str, Any], layer_index: int, vie
                 attention_output = attention_outputs[module_name]['output']
                 if isinstance(attention_output, list) and len(attention_output) >= 2:
                     # Get attention weights (element 1 of the output tuple)
-                    attention_weights = torch.tensor(attention_output[1])  # [batch, heads, seq, seq]
+                    raw_weights = attention_output[1]
+                    if raw_weights is None:
+                        continue  # Skip layers with missing attention data
+                    attention_weights = torch.tensor(raw_weights)  # [batch, heads, seq, seq]
                     layer_attention_pairs.append((layer_num, attention_weights))
         
         if not layer_attention_pairs: