Hello! I have a use case where during the training of a hybrid quantum/classical model at the end of an epoch I would like to change the measurement that occurs at the end of a quantum circuit. Here is a basic hybrid model to exemplify.
import pennylane as qml
import torch
n_qubits = 2
dev = qml.device("default.qubit", wires=n_qubits)
@qml.qnode(dev)
def qnode(inputs, weights):
qml.AngleEmbedding(inputs, wires=range(n_qubits))
qml.BasicEntanglerLayers(weights, wires=range(n_qubits))
return qml.sample()
n_layers = 2
weight_shapes = {"weights": (n_layers, n_qubits)}
qlayer = qml.qnn.TorchLayer(qnode, weight_shapes)
clayer_1 = torch.nn.Linear(2, 2)
clayer_2 = torch.nn.Linear(2, 2)
softmax = torch.nn.Softmax(dim=1)
layers = [clayer_1, qlayer, clayer_2, softmax]
model = torch.nn.Sequential(*layers)
During the training of this model (with some loss function/data) at the end of the epoch I would instead like to measure the probability of each computational basis using qml.probs(). Ideally, the function would behave as the following I believe:

Accepts the Qnode as an argument

Replaces the measurement at the end to qml.probs()

Executes the circuit to return the [len(2**n_qubits] probability list

Replaces this measurement with the original qml.sample() so at the beginning of the next epoch, training continues as previous
Any help would be greatly appreciated, whether using a method like this or something simpler. It seems that once you’ve compiled the Qnode, accessing the operations and measurements externally is difficult to do.