fix: handle cases in which losses are nan during the training

fix: handle cases in which losses are nan during the training #1464