Adamax Gradient Optimizer Update Step - Coding problem | OneNoughtOne