pytorch loss decrease slow