Fix a bug in train correctness checking
Summary: The train correctness checking requires the model to first run in eval mode. If the correctness checking fails, we should set the mode back to train before returning False.
Reviewed By: brad-mengchi
Differential Revision: D43083931
fbshipit-source-id: ba330bd6749e0e5043a7c83028a64e13a1ce412e
