Solved by ptrblck in post #4 
              
              
                No, setting the model to .eval() will change the behavior of some layers (e.g. dropout will be disabled and batchnorm layers will use their running stats) and will not change the behavior of the gradient calculation. 

Yes, that works.