Why is the size of the NTK matrices set to 300x300? Also the code doesnt work on tensorflow 2 I think. because of the compute_ntk function reshapes.