Incorrect shape assertion in test_decode_prefixes_device_handling

Nobody has claimed this yet.

## Bug Description

The test `test_decode_prefixes_device_handling` in `tests/test_nn_objectives.py` has an incorrect shape assertion on line 491.

## Current Code
```python
f_x = torch.randn(2, 32).cuda()
prefixes = torch.tensor([8, 16, 32])
x_hats = sae.decode(f_x, prefixes=prefixes)
assert x_hats.shape == (3, 2, 16) # BUG: Wrong shape!
```

## Expected Behavior

The `decode` method returns tensors with shape `(batch, n_prefixes, d_model)`. Given:
- batch_size = 2
- n_prefixes = 3
- d_model = 16

The correct shape should be `(2, 3, 16)` not `(3, 2, 16)`.

## Impact

This test is currently skipped on systems without CUDA, so the bug hasn't been caught by CI. However, on CUDA-enabled systems, this test would fail.

## Fix

Change line 491 from:
```python
assert x_hats.shape == (3, 2, 16)
```

to:
```python
assert x_hats.shape == (2, 3, 16)
```

Open tests/test_nn_objectives.py and inspect test_decode_prefixes_device_handling around line 491, including the expected decode shape. Run the test on a CUDA-enabled system and update the assertion so it matches the documented (batch, n_prefixes, d_model) order.

Written by the indexing model from the issue text.