diff --git a/docs/source/user_guide/feed_forward.rst b/docs/source/user_guide/feed_forward.rst
index 3c8cff7a..4cb889b2 100644
--- a/docs/source/user_guide/feed_forward.rst
+++ b/docs/source/user_guide/feed_forward.rst
@@ -124,6 +124,231 @@ without passing a tensor (an empty feature tensor is injected automatically).
    are not accepted as ``input_state``; wrap amplitude tensors with
    :meth:`~merlin.core.state_vector.StateVector.from_tensor` first.
 
+Known Limitation: Input Parameters Used Only Inside a Branch
+--------------------------------------------------------------
+
+.. note::
+
+   This is a known limitation of ``FeedForwardBlock`` in MerLin 0.4, tracked in
+   `issue #274 <https://github.com/merlinquantum/merlin/issues/274>`_. It is
+   expected to be solved in **MerLin 0.5** with a new ``FeedForwardBlock`` backend.
+   Until then, we propose the manual workaround below.
+
+``FeedForwardBlock`` currently requires that every entry of ``input_parameters`` be
+consumed by the **first** pre-measurement stage of the experiment. If a classical
+parameter is only referenced inside one or more ``pcvl.FFCircuitProvider`` branch
+configurations — i.e. it encodes a value *after* the measurement, in a specific
+branch — construction fails immediately with:
+
+.. code-block:: text
+
+   ValueError: The first stage must use all of the input parameters. Create you own
+   stages with variable input parameters with the partial measurement strategy instead
+
+For example, the following experiment is rejected because ``x`` is only used inside the
+branch circuits of the ``FFCircuitProvider``, not in the prefix unitary:
+
+.. code-block:: python
+
+   import perceval as pcvl
+   from perceval import BasicState, Circuit
+   from merlin.algorithms.feed_forward import FeedForwardBlock
+
+   # Build a 4-mode experiment with a prefix unitary and measurement in mode 0
+   experiment = pcvl.Experiment(4)
+   experiment.add(0, Circuit(4) // pcvl.Unitary.random(4))  # prefix unitary
+   experiment.add(0, pcvl.Detector.pnr())  # measure mode 0
+
+   # Branch-local parameter "x" is only defined inside branch circuits
+   provider = pcvl.FFCircuitProvider(1, 0, Circuit(3))
+   provider.add_configuration([0], pcvl.BS(pcvl.P("x")) // pcvl.BS(pcvl.P("A")))
+   provider.add_configuration([1], pcvl.BS(pcvl.P("x")) // pcvl.BS(pcvl.P("B")))
+   experiment.add(0, provider)
+
+   # FeedForwardBlock rejects this: "x" is not consumed by the first (prefix) stage
+   FeedForwardBlock(
+       experiment,
+       input_state=BasicState([1, 1, 0, 0]),
+       trainable_parameters=["A", "B"],
+       input_parameters=["x"],   # raises ValueError: "x" is not used in the first stage
+   )
+
+**Workaround: build the stages manually with partial measurement**
+
+Until MerLin 0.5 ships, the same physical workflow can be reproduced by chaining
+:class:`~merlin.algorithms.layer.QuantumLayer` instances yourself with the
+:doc:`partial measurement strategy </quantum_expert_area/partial_measurement>`.
+Critically, to make the model trainable, you must construct layers **once** in
+an ``nn.Module.__init__``, then call :meth:`~merlin.algorithms.layer.QuantumLayer.set_input_state`
+at runtime to route conditional amplitudes through them. This is the only safe way
+to preserve trainable parameters across training steps.
+
+The pattern closely mirrors what ``FeedForwardBlock`` does internally for its
+supported case (branch-local inputs aside). However, this workaround covers a
+**single feedforward stage**; multi-stage experiments require nesting the same
+pattern per stage.
+
+Key implementation details:
+
+- **Layer construction:** Build the prefix ``QuantumLayer`` and one branch layer
+  per outcome in ``__init__``, stored in an ``nn.ModuleDict`` (or similar container).
+- **Runtime routing:** In ``forward()``, call ``set_input_state(branch.amplitudes)``
+  on each branch layer to inject the conditional measurement outcome.
+- **Batch handling:** The prefix layer is called once with no features (produces
+  batch size 1), while branch layers receive input with shape ``(batch_size, ...)``.
+  Probabilities must be broadcast correctly: use ``unsqueeze(-1)`` or indexing to
+  ensure branch and conditional probabilities are multiplied element-wise.
+
+.. code-block:: python
+
+   import torch
+   import torch.nn as nn
+   import perceval as pcvl
+   from perceval import Circuit
+   from merlin.algorithms.layer import QuantumLayer
+   from merlin.core.computation_space import ComputationSpace
+   from merlin.measurement.strategies import MeasurementStrategy
+
+   class BranchFeedforward(nn.Module):
+       """Feedforward model with branch-local parameters, trainable via gradient descent.
+
+       **Note:** This example assumes a single measured mode (e.g., measured_modes=[0]).
+       The implementation below reconstructs full-system keys for arbitrary
+       measured modes.
+       """
+
+       def __init__(self, input_state, prefix_circuit):
+           super().__init__()
+
+           # Build the prefix stage (partial measurement after mode 0).
+           self.partial_layer = QuantumLayer(
+               circuit=prefix_circuit,
+               input_state=input_state,
+               measurement_strategy=MeasurementStrategy.partial(
+                   modes=[0],
+                   computation_space=ComputationSpace.FOCK,
+               ),
+           )
+
+           # Pre-build branch layers for each outcome, keyed by measurement outcome.
+           # Initialize with input_size (classical) and n_photons (quantum dimension).
+           # A branch layer uses the photon count remaining after its measured
+           # outcome, matching the routed conditional state dimension.
+           self.branch_layers = nn.ModuleDict()
+
+           # Build 2-mode branch circuits (remaining modes after measuring mode 0).
+           c0 = Circuit(2)
+           c0.add(0, pcvl.BS(pcvl.P("x")))
+           c0.add(0, pcvl.BS(pcvl.P("A")))
+
+           c1 = Circuit(2)
+           c1.add(0, pcvl.BS(pcvl.P("x")))
+           c1.add(0, pcvl.BS(pcvl.P("B")))
+
+           c2 = Circuit(2)  # vacuum branch after measuring two photons
+
+           self.branch_configs = {
+               (0,): (c0, ["A"], ["x"]),
+               (1,): (c1, ["B"], ["x"]),
+               (2,): (c2, [], []),
+           }
+
+           for outcome, (circuit, trainable_params, input_params) in self.branch_configs.items():
+               key = str(outcome)
+               if sum(input_state) - sum(outcome) == 0:
+                   continue
+               self.branch_layers[key] = QuantumLayer(
+                   circuit=circuit,
+                   input_size=len(input_params),  # Classical input dimension
+                   n_photons=sum(input_state) - sum(outcome),
+                   trainable_parameters=trainable_params,
+                   input_parameters=input_params,
+                   measurement_strategy=MeasurementStrategy.probs(ComputationSpace.FOCK),
+               )
+
+       def forward(self, x=None):
+           # Run the prefix stage to get partial measurement branches.
+           # Note: called with no features; produces batch size 1.
+           partial_measurement = self.partial_layer()
+
+           probabilities = {}
+           for branch in partial_measurement.branches:
+               key = str(branch.outcome)
+               if sum(branch.outcome) == sum(self.partial_layer.input_state):
+                   vacuum_key = [0] * len(self.partial_layer.input_state)
+                   for mode, value in zip(partial_measurement.measured_modes, branch.outcome):
+                       vacuum_key[mode] = value
+                   probabilities[tuple(vacuum_key)] = branch.probability.expand(
+                       x.shape[0] if x is not None else 1
+                   )
+                   continue
+               branch_layer = self.branch_layers[key]
+
+               # Set the conditional amplitudes for this branch.
+               # set_input_state() is the only way to route a StateVector through
+               # a layer that was already constructed with trainable parameters.
+               branch_layer.set_input_state(branch.amplitudes)
+
+               # Execute the branch layer with its local classical inputs (if any).
+               _circuit, _trainable_params, input_params = self.branch_configs[branch.outcome]
+               if input_params:
+                   branch_output = branch_layer(x)  # shape: (batch_size, n_keys)
+               else:
+                   branch_output = branch_layer()  # shape: (batch_size, n_keys)
+
+               # Weight by branch probability with proper broadcasting.
+               # branch.probability is (batch_size,), branch_output is (batch_size, n_keys).
+               # unsqueeze(-1) broadcasts branch probability to (batch_size, 1).
+               branch_prob_weighted = branch.probability.unsqueeze(-1)  # (batch_size, 1)
+
+               for index, remaining_key in enumerate(branch_layer.output_keys):
+                   # Reconstruct the full-system key in the original mode order.
+                   output_key = [0] * len(self.partial_layer.input_state)
+                   for mode, value in zip(partial_measurement.measured_modes, branch.outcome):
+                       output_key[mode] = value
+                   for mode, value in zip(partial_measurement.unmeasured_modes, remaining_key):
+                       output_key[mode] = value
+                   output_key = tuple(output_key)
+                   branch_probs_for_key = branch_output[:, index]  # (batch_size,)
+                   weighted_probs = branch_prob_weighted.squeeze(-1) * branch_probs_for_key
+                   probabilities[output_key] = weighted_probs
+
+           return probabilities
+
+   # Usage
+   input_state = [1, 1, 0]
+   prefix = Circuit(3) // pcvl.Unitary.random(3)
+   model = BranchFeedforward(input_state, prefix)
+
+   # Single batch
+   x = torch.tensor([[0.2]])  # (batch_size=1, input_dim=1)
+   probs_single = model(x)
+
+   # Multi-batch (batch_size=3)
+   x = torch.tensor([[0.1], [0.2], [0.3]])  # (batch_size=3, input_dim=1)
+   probs_multi = model(x)  # each probability has shape (3,)
+
+   # Verify probabilities sum to 1 per batch element
+   for i in range(x.shape[0]):
+       batch_total = sum(p[i].item() for p in probs_multi.values())
+       assert abs(batch_total - 1.0) < 1e-5
+
+   # Now the model is trainable:
+   optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+   # ... training loop ...
+
+This manual composition reproduces the same output structure as ``FeedForwardBlock``
+for single-stage PNR experiments. It uses partial Fock measurement rather than
+``FeedForwardBlock``'s amplitude-based detector transform, so equivalence should
+not be assumed for non-PNR detectors or multi-stage experiments.
+
+.. note::
+
+   The pattern is verified by the test suite in
+   ``tests/algorithms/test_feedforward_manual_workaround.py``, which covers
+   single-batch, multi-batch, trainability, and error cases. When generalizing
+   to multiple measured modes (e.g., ``measured_modes=[0, 1]``), ensure key
+   reconstruction uses the correct measured modes and remaining modes.
 
 Further Reading
 ---------------
diff --git a/tests/algorithms/test_feedforward_manual_workaround.py b/tests/algorithms/test_feedforward_manual_workaround.py
new file mode 100644
index 00000000..4a1db6b2
--- /dev/null
+++ b/tests/algorithms/test_feedforward_manual_workaround.py
@@ -0,0 +1,357 @@
+# MIT License
+#
+# Copyright (c)
+#
+# Tests for the manual feedforward workaround using partial measurement and set_input_state.
+# This verifies the pattern documented in docs/source/user_guide/feed_forward.rst
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import perceval as pcvl
+import pytest
+from perceval import BasicState, Circuit
+
+from merlin.algorithms.layer import QuantumLayer
+from merlin.algorithms.feed_forward import FeedForwardBlock
+from merlin.core.computation_space import ComputationSpace
+from merlin.measurement.strategies import MeasurementStrategy
+
+
+class BranchFeedforward(nn.Module):
+    """Feedforward model with branch-local parameters, trainable via gradient descent.
+    
+    This is the pattern documented in docs/source/user_guide/feed_forward.rst
+    for the workaround to FeedForwardBlock's limitation on branch-local parameters.
+    
+    Parameters
+    ----------
+    input_state : list | tuple
+        Initial Fock state for the prefix circuit.
+    prefix_circuit : Circuit
+        Perceval circuit for the prefix (pre-measurement) stage.
+    branch_configs : dict
+        Mapping from measurement outcome (tuple) to (circuit, trainable_params, input_params).
+    measured_modes : list[int]
+        Modes to measure in the partial measurement (e.g., [0]).
+    """
+
+    def __init__(self, input_state, prefix_circuit, branch_configs, measured_modes):
+        super().__init__()
+        self.measured_modes = measured_modes
+        self.branch_configs_dict = branch_configs
+
+        # Build the prefix stage (partial measurement).
+        self.partial_layer = QuantumLayer(
+            circuit=prefix_circuit,
+            input_state=input_state,
+            measurement_strategy=MeasurementStrategy.partial(
+                modes=measured_modes,
+                computation_space=ComputationSpace.FOCK,
+            ),
+        )
+
+        # Pre-build branch layers for each outcome, keyed by measurement outcome.
+        # Initialize with input_size (classical) and n_photons (quantum dimension).
+        # The branch layer uses the photon count remaining after measurement.
+        self.branch_layers = nn.ModuleDict()
+        
+        for outcome, (circuit, trainable_params, input_params) in branch_configs.items():
+            key = str(outcome)
+            remaining_photons = sum(input_state) - sum(outcome)
+            if remaining_photons == 0:
+                continue
+            self.branch_layers[key] = QuantumLayer(
+                circuit=circuit,
+                input_size=len(input_params),  # Classical input dimension
+                n_photons=remaining_photons,
+                trainable_parameters=trainable_params,
+                input_parameters=input_params,
+                measurement_strategy=MeasurementStrategy.probs(ComputationSpace.FOCK),
+            )
+
+    def forward(self, x: torch.Tensor | None = None) -> dict[tuple, torch.Tensor]:
+        """Forward pass through the branched model.
+        
+        Parameters
+        ----------
+        x : torch.Tensor | None
+            Classical input tensor for branches that require input_parameters.
+            If provided, shape should be (batch_size, input_dim).
+            If None, used for branches with no input_parameters.
+        
+        Returns
+        -------
+        dict[tuple, torch.Tensor]
+            Dictionary mapping full measurement outcome keys to weighted probabilities.
+            Probabilities have shape (batch_size,) and sum to 1 over all keys.
+        """
+        # Run the prefix stage to get partial measurement branches.
+        partial_measurement = self.partial_layer()
+
+        probabilities = {}
+        measured_modes = partial_measurement.measured_modes
+        unmeasured_modes = partial_measurement.unmeasured_modes
+
+        def reconstruct_full_key(remaining_key: tuple[int, ...]) -> tuple[int, ...]:
+            full_key = [0] * len(self.partial_layer.input_state)
+            for mode, value in zip(measured_modes, branch.outcome):
+                full_key[mode] = value
+            for mode, value in zip(unmeasured_modes, remaining_key):
+                full_key[mode] = value
+            return tuple(full_key)
+
+        for branch in partial_measurement.branches:
+            key = str(branch.outcome)
+            if sum(branch.outcome) == sum(self.partial_layer.input_state):
+                probabilities[reconstruct_full_key((0,) * len(unmeasured_modes))] = branch.probability.expand(
+                    x.shape[0] if x is not None else 1
+                )
+                continue
+            branch_layer = self.branch_layers[key]
+
+            # Set the conditional amplitudes for this branch.
+            # set_input_state() is the only way to route a StateVector through
+            # a layer that was already constructed with trainable parameters.
+            branch_layer.set_input_state(branch.amplitudes)
+
+            # Execute the branch layer with its local classical inputs (if any).
+            circuit, _trainable_params, input_params = self.branch_configs_dict[branch.outcome]
+            if input_params:
+                branch_output = branch_layer(x)  # shape: (batch, n_keys)
+            else:
+                branch_output = branch_layer()  # shape: (batch, n_keys)
+
+            # Weight by branch probability and store results.
+            # branch.probability is (batch,), branch_output is (batch, n_keys).
+            # For each outcome key, weight by branch probability.
+            branch_prob_weighted = branch.probability.unsqueeze(-1)  # (batch, 1)
+
+            for index, remaining_key in enumerate(branch_layer.output_keys):
+                output_key = reconstruct_full_key(remaining_key)
+                branch_probs_for_key = branch_output[:, index]  # (batch,)
+                weighted_probs = branch_prob_weighted.squeeze(-1) * branch_probs_for_key
+                probabilities[output_key] = weighted_probs
+
+        return probabilities
+
+
+def test_manual_feedforward_workaround_single_batch():
+    """Verify the manual workaround works for batch_size=1."""
+    input_state = [1, 1, 0]
+    prefix = Circuit(3) // pcvl.Unitary.random(3)
+
+    # Define branch-specific circuits explicitly as 2-mode circuits (after measuring mode 0).
+    # Use explicit Circuit(2) to ensure proper dimensionality.
+    c0 = Circuit(2)
+    c0.add(0, pcvl.BS(pcvl.P("x")))
+    c0.add(0, pcvl.BS(pcvl.P("A")))
+    
+    c1 = Circuit(2)
+    c1.add(0, pcvl.BS(pcvl.P("x")))
+    c1.add(0, pcvl.BS(pcvl.P("B")))
+
+    c2 = Circuit(2)
+    
+    branch_configs = {
+        (0,): (c0, ["A"], ["x"]),
+        (1,): (c1, ["B"], ["x"]),
+        (2,): (c2, [], []),
+    }
+
+    model = BranchFeedforward(input_state, prefix, branch_configs, measured_modes=[0])
+
+    # Single batch forward pass
+    x = torch.tensor([[0.2]])  # (batch=1, input_dim=1)
+    probabilities = model(x)
+
+    # Verify output structure
+    assert isinstance(probabilities, dict), "Output should be a dictionary"
+    assert len(probabilities) > 0, "Should have at least one probability entry"
+
+    # Verify all keys are tuples and all values are tensors with correct shape
+    for key, prob in probabilities.items():
+        assert isinstance(key, tuple), f"Key {key} should be a tuple"
+        assert isinstance(prob, torch.Tensor), f"Probability for {key} should be a tensor"
+        assert prob.shape == (1,), f"Probability for batch_size=1 should have shape (1,), got {prob.shape}"
+        assert 0 <= prob.item() <= 1, f"Probability {prob.item()} should be in [0, 1]"
+
+    # Verify probabilities sum to 1
+    total_prob = sum(p.item() for p in probabilities.values())
+    assert abs(total_prob - 1.0) < 1e-5, f"Probabilities should sum to 1, got {total_prob}"
+
+
+def test_manual_feedforward_workaround_multi_batch():
+    """Verify the manual workaround works correctly for batch_size > 1 with proper broadcasting."""
+    input_state = [1, 1, 0]
+    prefix = Circuit(3) // pcvl.Unitary.random(3)
+
+    # Define branch-specific circuits explicitly as 2-mode circuits.
+    c0 = Circuit(2)
+    c0.add(0, pcvl.BS(pcvl.P("x")))
+    c0.add(0, pcvl.BS(pcvl.P("A")))
+    
+    c1 = Circuit(2)
+    c1.add(0, pcvl.BS(pcvl.P("x")))
+    c1.add(0, pcvl.BS(pcvl.P("B")))
+
+    c2 = Circuit(2)
+    
+    branch_configs = {
+        (0,): (c0, ["A"], ["x"]),
+        (1,): (c1, ["B"], ["x"]),
+        (2,): (c2, [], []),
+    }
+
+    model = BranchFeedforward(input_state, prefix, branch_configs, measured_modes=[0])
+
+    # Multi-batch forward pass
+    batch_size = 3
+    x = torch.tensor([[0.1], [0.2], [0.3]])  # (batch=3, input_dim=1)
+    probabilities = model(x)
+
+    # Verify all probabilities have the correct batch dimension
+    for key, prob in probabilities.items():
+        assert prob.shape == (batch_size,), (
+            f"Probability for {key} should have shape ({batch_size},), got {prob.shape}"
+        )
+        assert torch.all((prob >= 0) & (prob <= 1)), (
+            f"All probabilities for {key} should be in [0, 1], got {prob}"
+        )
+
+    # Verify probabilities sum to 1 for each batch element independently
+    prob_array = torch.stack(list(probabilities.values()), dim=0)  # (n_keys, batch_size)
+    total_probs = prob_array.sum(dim=0)  # (batch_size,)
+    for i, total in enumerate(total_probs):
+        assert abs(total.item() - 1.0) < 1e-5, (
+            f"Batch element {i}: probabilities sum to {total.item()}, expected 1.0"
+        )
+
+
+def test_manual_feedforward_workaround_trainable():
+    """Verify that parameters remain trainable across multiple forward passes."""
+    input_state = [1, 1, 0]
+    prefix = Circuit(3) // pcvl.Unitary.random(3)
+
+    # Define branch-specific circuits explicitly as 2-mode circuits.
+    c0 = Circuit(2)
+    c0.add(0, pcvl.BS(pcvl.P("x")))
+    c0.add(0, pcvl.BS(pcvl.P("A")))
+    
+    c1 = Circuit(2)
+    c1.add(0, pcvl.BS(pcvl.P("x")))
+    c1.add(0, pcvl.BS(pcvl.P("B")))
+
+    c2 = Circuit(2)
+    
+    branch_configs = {
+        (0,): (c0, ["A"], ["x"]),
+        (1,): (c1, ["B"], ["x"]),
+        (2,): (c2, [], []),
+    }
+
+    model = BranchFeedforward(input_state, prefix, branch_configs, measured_modes=[0])
+
+    # Collect initial parameter values
+    initial_params = {
+        name: param.clone().detach()
+        for name, param in model.named_parameters()
+    }
+
+    # Define a simple loss and optimize
+    x = torch.tensor([[0.2]])
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+    for _ in range(5):
+        optimizer.zero_grad()
+        probs = model(x)
+        # Simple loss: try to maximize the first probability we encounter
+        first_probability_by_branch = {
+            key[0]: probability
+            for key, probability in probs.items()
+            if key[0] in (0, 1)
+        }
+        loss = -sum(
+            probability.mean() for probability in first_probability_by_branch.values()
+        )
+        loss.backward()
+        optimizer.step()
+
+    # Verify parameters have changed
+    for name, initial_param in initial_params.items():
+        current_param = dict(model.named_parameters())[name]
+        assert not torch.allclose(initial_param, current_param, atol=1e-6), (
+            f"Parameter {name} should have been updated by optimizer"
+        )
+
+
+def test_feedforwardblock_input_at_branch_fails():
+    """Verify that FeedForwardBlock raises ValueError for branch-local input parameters.
+    
+    This is the error case that the manual workaround solves.
+    This test reproduces the scenario from the documentation error example.
+    """
+    # Build a complete experiment with prefix, measurement, and branch-local parameter
+    m = 4
+    prefix = Circuit(m) // pcvl.Unitary.random(m)
+
+    # Branch-local trainable parameters are distinct from the shared input.
+    experiment = pcvl.Experiment(m)
+    experiment.add(0, prefix)
+    experiment.add(0, pcvl.Detector.pnr())  # Measure first mode
+
+    # FFCircuitProvider with branch-local parameters
+    # Use unique parameter names per branch to avoid circuit conflicts
+    x = pcvl.P("x")
+    c0 = Circuit(m - 1)
+    c0.add(0, pcvl.BS(x) // pcvl.BS(pcvl.P("A")))
+    
+    c1 = Circuit(m - 1)
+    c1.add(0, pcvl.BS(x) // pcvl.BS(pcvl.P("B")))
+    
+    provider = pcvl.FFCircuitProvider(1, 0, Circuit(m - 1))
+    provider.add_configuration([0], c0)
+    provider.add_configuration([1], c1)
+    experiment.add(0, provider)
+
+    # FeedForwardBlock should reject this because x is only used in branches.
+    with pytest.raises(ValueError, match="The first stage must use all of the input parameters"):
+        FeedForwardBlock(
+            experiment,
+            input_state=BasicState([1, 1, 0, 0]),
+            trainable_parameters=["A", "B"],
+            input_parameters=["x"],
+        )
+
+
+def test_manual_feedforward_keys_arbitrary_measured_modes():
+    """Verify key reconstruction places outcomes in their measured modes."""
+    input_state = [1, 1, 0]
+    prefix = Circuit(3) // pcvl.Unitary.random(3)
+
+    # Define branch-specific circuits explicitly as 2-mode circuits.
+    c0 = Circuit(2)
+    c0.add(0, pcvl.BS(pcvl.P("x")))
+    c0.add(0, pcvl.BS(pcvl.P("A")))
+    
+    c1 = Circuit(2)
+    c1.add(0, pcvl.BS(pcvl.P("x")))
+    c1.add(0, pcvl.BS(pcvl.P("B")))
+
+    c2 = Circuit(2)
+    
+    branch_configs = {
+        (0,): (c0, ["A"], ["x"]),
+        (1,): (c1, ["B"], ["x"]),
+        (2,): (c2, [], []),
+    }
+
+    model = BranchFeedforward(input_state, prefix, branch_configs, measured_modes=[1])
+
+    x = torch.tensor([[0.2]])
+    probabilities = model(x)
+
+    assert probabilities
+    for key in probabilities:
+        assert key[1] in (0, 1, 2)