In this function in `model.py`, you divide by the total length of the dimension by the shape of the dimension, instead of the "shape-1".