Hello there !

First, I am sorry if my post does not fit exactly the autograd category but it seemed to be still a little bit related to it.

Alright I’ll try to explain my issue as clearly as possible !

Long story short : not all the weights in my modules are updated.

The detailed story :

I have a first network returning 2 values : u_lin and v – very classic feedforward static network.

```
u_lin, v = u_lin_mod(batch_u, batch_x)
```

u_lin and v are produced by alpha and beta layers.

It is where trouble begins.

So I am dealing with “Neural ODEs” so there is module which is an integration scheme (RK4).

Then v is fed into this RK4 module which is basically a fancy for loop :

```
z_sim_torch_fit, y_sim_torch_fit = M2_solution(v, batch_z0_hidden)
```

z0 is the initial condition for integration.

The M2_solution constructor :

```
def __init__(self, ss_model, ts=1.0, scheme='RK44', device="cpu"):
super(RK4Simulator, self).__init__()
self.ss_model = ss_model
self.ts = ts
self.device = device
```

The forward method :

```
def forward(self, u_batch, x0_batch):
""" Multi-step simulation over (mini)batches
Parameters
----------
u_batch: Tensor. Size: (m, q, n_u)
Input sequence for each subsequence in the minibatch
x0_batch: Tensor. Size: (q, n_x)
initial state for each sequence in the minibatch
x_batch: Tensor. Size: (m, q, n_x)
state sequence for each subsequence in the minibatch
Returns
-------
Tensor. Size: (m, q, n_x)
Simulated state for all subsequences in the minibatch
"""
X_sim_list = []
Y_sim_list: List[torch.Tensor] = []
x_step = x0_batch
for u_step in u_batch.split(1):#i in range(seq_len):
u_step = u_step.squeeze(0)
# x_step = x_step.squeeze(0)
X_sim_list += [x_step]
dt2 = self.ts / 2.0
k1,y_step = self.ss_model(u_step, x_step)
Y_sim_list += [y_step]
k2_dx, _ = self.ss_model(u_step, x_step + dt2*k1)
k3_dx, _ = self.ss_model(u_step, x_step + dt2*k2_dx)
k4_dx, _ = self.ss_model(u_step, x_step + self.ts*k3_dx)
dx = self.ts / 6.0 * (k1 + 2.0 * k2_dx + 2.0 * k3_dx + k4_dx)
x_step = x_step + dx
X_sim = torch.stack(X_sim_list, 0)
Y_sim = torch.stack(Y_sim_list, 0)
return X_sim, Y_sim
```

So when I do `loss.backward()`

only the weights of M2_solution are updated and not the ones of u_lin_mod.

However I want the RK4 class to be as generic as possible and I do not want to call it with u_batch and x_batch.

So I tried to change the constructor to this :

```
def __init__(self, ss_model, u_lin, ts=1.0, scheme='RK44', device="cpu"):
super(RK4Simulator, self).__init__()
self.ss_model = ss_model
self.ts = ts
self.device = device
self.u_lin_model = u_lin
```

and in the main loop to call :

```
u_lin, v = M2_solution.u_lin_model(batch_u, batch_x)
z_sim_torch_fit, y_sim_torch_fit = M2_solution(v, batch_z0_hidden)
```

but it does not work either. Is there a way to work around without writing a lot of different classes. (Because I have 2 other group of modules performing the same operations).

Thanks a lot guys !