* change learning_rate_decay to learning_rate_scheduler
* optimize code
* change nn.cast to tensor.cast