@mrshenli Thanks for your suggestion! I found an interesting behavior. Previously the model was trained on CPU and the training procedure generated these messages. But once the training is moved to GPU, using find_unused_parameters=True works totally fine without any error or warning.