I tried both your methods and they worked for me. My guess is that your bug is in the way you compute the Gram Matrices: it does not work if you simply reshape the a_S and a_G to the shapes you want, because that scrambles the data. Here’s a thread which explains why. You may need to read some of the earlier posts on that thread to get the full context.
And here’s a thread that explains why the integers need to be cast to floats depending on how you write the code.