@mr_fnord, did you include axis=1 in your loss function? I was tipped off to that by this post. I was getting the same error as you before I did that.
axis=1