I want to rebuild this machine learning structure, a KalmanNet. Machine Learning Structure KalmanNet. I have problems with the pointwise multiplication.
This is my code:

class MB_KalmanNet_GRU(torch.nn.Module):
    def __init__(self, F, in_gru, hidden_gru):
        super().__init__()
        self.F = F

        self.FC11 = nn.Linear(6, in_gru).to(device)
        self.GRU1 = nn.GRU(in_gru, hidden_gru).to(device)
        self.FC21 = nn.Linear(hidden_gru,  9).to(device)

    def forward(self, est, h_0, track_first):
        track_total = torch.empty(est.shape).to(device)
        track_total[0, :,: ] = track_first[0, :, :]
        pred_prev = track_first[0, :, :]
        h = h_0.clone()

        for i in range(1, est.shape[0]):
            pred = track_total[i-1, :,: ].unsqueeze(0) @ self.F.T

            delta_y = est[i, :, :] - pred
            delta_x = pred_prev - track_total[i-1, :,: ].unsqueeze(0)
            pred_prev = pred

            input_KGain = torch.cat((delta_y, delta_x), dim=2)
            into_gru = self.FC11(input_KGain)
            outta_gru, h = self.GRU1(into_gru, h)
            KGain = self.FC21(outta_gru)

            KGain_matrix = torch.reshape(KGain, (est.shape[1], 3, 3))
            innovation = torch.bmm(KGain_matrix, delta_y.squeeze(0).unsqueeze(2))
           
            track_total[i, :, :] = pred[0, :, :] + innovation.squeeze(2).unsqueeze(0)

        return track_total

I have problems with the section

           innovation = torch.bmm(KGain_matrix, delta_y.squeeze(0).unsqueeze(2))

If delta_y is just a tensor the gradiants are calculated correctly. Also if I add KGain_matrix and delta_y.squeeze(0).unsqueeze(2) everything works fine. Only whenmultiplication is performed the gradiants are not correctly calculated anymore. Thank you so much in advance.

I am expecting that the gradiants are calculated correctly, but instead they are NAN. I tried differnt structures and they work, just the multiplication which i need doesent work.