Needed to add Wasserstain/Hinge loss instead of BCE
Needed to add Wasserstain/Hinge loss instead of BCE