I’m still skeptical about your solution of making sure the target tensor is attached to a computation graph as mentioned here, since it’s common to assume a static target while the model outputs should of course be attached to the computation graph created in the forward pass of the model.