Train different model after each epoch for evaluation

I am trying to train a model that generates embeddings. I want to evaluate this model every epoch based on its output (embeddings). So I create a simple DNN which takes these embeddings as input and performs classification (also trains several epochs every evaluation step).

However, it seems PyTorch is not able to train a decoupled model during the training loop itself. I am always getting an RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn error when training the simple DNN for evaluation. Training the DNN decoupled of the outer training loop for the embeddings works fine!

Is there a trick to train 2 models simultaneously, but with different computational graphs?


Found a different solution.