Hi,
thanks a lot for the solution. I was not aware of the momentum buffers. And I haven’t specially treated them in my previous work. Why do you think this happened? To my knowledge, I only need to push the model and loss function to current device. Could you describe the best practice in this case please?