Two models with .eval() and .train() separately but always incompatible

Hi,

I have two models with .eval() (model_A) and .train() (model_B) separately. I am inputting the output of model_A, a trained model, to model_B which I am training, as followed.

model_A.eval()
for epoch in range(epoch_num):

model_B.train()
output_model_A = model_A(input)
output_mobel_B = model_B(output_model_A)

batch_loss.backward()
optimizer.step()

It popped up an error “RuntimeError: cudnn RNN backward can only be called in training mode”. I have set the model_B to train mode. I don’t know why I still get this kind of error.

Does anyone have a solution? Many thanks!!!

Wrap model_A with torch.no_grad() or use .detach() on the outputs(preferably the former for avoiding memory issues).

with torch.no_grad():
    data=model_A(input)

#rest of training not wrapped
1 Like