Hi @shangguan91, the code is correct.
Calling model.train() doesn’t train your model, it just sets your model in training mode, so virtually no cost in time.
For more information, please refer to here 2.