Parametrized Quantum Circuits for Reinforcement Learning stuck if using qsim backend

I've ran this notebook with no problem. However, if I try to use qsim as a backend by passing the backend backend = qsimcirq.QSimSimulator() explicitly to ReUploadingPQC, then I get the following message

2023-04-26 09:20:03.862016: W tensorflow/core/grappler/optimizers/loop_optimizer.cc:907] Skipping loop optimization for Merge node with control input: cond/branch_executed/_12

and it gets stuck there indefinitely.

Any idea of what might be going on?

To reproduce:

replace

def generate_model_policy(
    qubits: List[cirq.GridQubit],
    n_layers: int,
    n_actions: int,
    beta: float,
    observables: List[cirq.PauliString],
) -> tf.keras.Model:
    """Generates a Keras model for a data re-uploading PQC policy."""

    input_tensor = tf.keras.Input(shape=(len(qubits),), dtype=tf.dtypes.float32, name="input")
    re_uploading_pqc = ReUploadingPQC(qubits, n_layers, observables)([input_tensor])
    process = tf.keras.Sequential(
        [Alternating(n_actions), tf.keras.layers.Lambda(lambda x: x * beta), tf.keras.layers.Softmax()],
        name="observables-policy",
    )
    policy = process(re_uploading_pqc)
    model = tf.keras.Model(inputs=[input_tensor], outputs=policy)

    return model

with

def generate_model_policy(
    qubits: List[cirq.GridQubit],
    n_layers: int,
    n_actions: int,
    beta: float,
    observables: List[cirq.PauliString],
    backend="noiseless",
) -> tf.keras.Model:
    """Generates a Keras model for a data re-uploading PQC policy."""

    input_tensor = tf.keras.Input(shape=(len(qubits),), dtype=tf.dtypes.float32, name="input")
    re_uploading_pqc = ReUploadingPQC(qubits, n_layers, observables, backend=backend)([input_tensor])
    process = tf.keras.Sequential(
        [Alternating(n_actions), tf.keras.layers.Lambda(lambda x: x * beta), tf.keras.layers.Softmax()],
        name="observables-policy",
    )
    policy = process(re_uploading_pqc)
    model = tf.keras.Model(inputs=[input_tensor], outputs=policy)

    return model

and then replace

model = generate_model_policy(qubits, n_layers, n_actions, 1.0, observables)

backend = qsimcirq.QSimSimulator()

model = generate_model_policy(qubits, n_layers, n_actions, 1.0, observables, backend=backend)

versions:

gym = "==0.22.0"
tensorflow = "2.7.0"
tensorflow-quantum = "0.7.2"
qsimcirq = "0.13.3" # <0.16.0 because of incompatibility with tensorflow-quantum