Further adjust test

Author: degenfabian

tests/integration/model_bridge/test_bridge_integration.py

@@ -318,7 +318,10 @@ def capture_pattern_hook(tensor, hook):
             len(pattern_tensor.shape) == 4
         ), f"Pattern tensor should be 4D, got {len(pattern_tensor.shape)}D"
 
-        n_heads_dim, pos_q_dim, pos_k_dim = pattern_tensor.shape
+        batch_dim, n_heads_dim, pos_q_dim, pos_k_dim = pattern_tensor.shape
+
+        # Verify the batch dimension is 1
+        assert batch_dim == 1, f"Batch dimension should be 1, got {batch_dim}"
 
         # Verify dimensions make sense
         assert (