https://github.com/vinhkhuc/PyTorch-Mini-Tutorials/blob/cbf48ed33aa32fb58bfee594e4b7a169bb95e283/2_logistic_regression.py#L47