Why the loss is nan by using linear activation function in the last layer?
There seem a lot of initialization problems.


There ain't too many biases. Its just 1 per layer. Your shape of
biases seems wrong.
You may have got the network right in shapes since it doesn't return shape errors. But bear in mind, it seems highly likely #1 is
true.
Same goes to w0 and w1.


Consider checking all dims again. If you need, I'm attaching a sample Neural Network implementation in python.

https://raw.githubusercontent.com/MasterSkepticista/NEAT_py/master/network.py

This is as it happens to be taught in Andrew Ng's course.