I am trying to implement a specific functionality using a neural network. I have attached the block diagram of my system. Basically, I am trying to develop a recurrent cell-like functionality where the output of a certain layer (Output of FC2)  is used as feedback input to a previous layer (CONV1).
The difference from an RNN cell is that the feedback should update the parameters of the CONV1 layer and the whole network’s parameters should be updated based on an intermediate layer loss.
A simplified version of the code is given below. The way I have implemented, FC2 layer parameters do not get updated during training.
class Net(nn.Module):
    def __init__(self, inputSize, num_classes):
        super(Net, self).__init__()

        # Filtering
        self.conv1 = custom_conv(cnn_N_filt[0], cnn_len_filt[0], fs)

        # Backend
        self.fc1 = nn.Linear(x, y) # Size based on previous layer output

        # Parameter calculation
        self.fc2 = nn.Linear(num_classes+inputSize, 1)

    def forward(self, x, hidden):
        inputTemp=x.view(-1,3200)
        x = F.relu(self.conv1(x, hidden))
        x = F.max_pool1d(x, 2, 2)
        x = x.view(-1, x.shape[1]*x.shape[2])
        x = F.relu(self.fc1(x))
        x = torch.cat((x, inputTemp), dim=1)
        hidden  = F.relu(self.fc2(x))
        hidden = hidden.view(-1)

        return F.log_softmax(x, dim=1), F.sigmoid(hidden)

So, I believe the hidden output from FC2 should be given as a model parameter to CONV1 during each forward pass.
Any ideas on how to achieve this?
Block_diagram.png2608×3317 126 KB