Kaiming Weight Initialization Strategy - Coding problem | OneNoughtOne