Replace nn.module call to train() for executorch models

Replace nn.module call to train() for executorch models #661