Mini Batch Gradient Descent Python