Hi, I’m considering to build a network like this:

- Each linear layer
**A**,**B**receive 128 dim output from**CNN**(Φ). -
**CNN**’s parameters are trained in another task, and saved as`model.pth`

. - Now I want to train
**A**,**B**layers and**CNN**in new task.

In this case, is next pseudo code works properly?

```
model = CNN() # class CNN() defines Φ's archtecture
model.load_state_dict(torch.load("model.pth"))
modelA = LinearUnitA()
modelB = LinearUnitB()
optim = optimizer_definition() # some operations to define optimizer
for i in N_training:
f_128 = model.forward(input_tensor)
outputA = modelA.forward(f_128)
outputB = modelA.forward(f_128)
lossA = loss_computation_A(outputA)
lossB = loss_computation_B(outputB)
lossA.backward()
lossB.backward()
optim.step()
```

Thanks.