I think this is a mistake, maybe affect the performance.

[Chienwang5]

https://github.com/snap-stanford/GreaseLM/blob/65d7ac84854293f630d6f8ce1428234126c74281/greaselm.py#L342C13-L342C33
It should be loss/=(b-a), not loss/=bs.