For numerical stability we can use the log-sum-exp trick to calculate some loss such as cross entropy