I think a better and elegant way to implement is the pytorch way. By defining a forward and backward function for easy chaining, you can easily chain multiple l

I think a better and elegant way to implement is the pytorch way. By defining a forward and backward function for easy chaining, you can easily chain multiple layers into a single model. Placing the training loop inside the layer will be hard to write a multi layer version. Just wrote a similar from scratch code few months back: https://github.com/theblackcat102/DL_2018/blob/master/HW2/layers.py https://github.com/theblackcat102/DL_2018/blob/master/HW2/la...