Training MNIST with MLX on Apple Silicon

import mlx.core as mx import mlx.nn as nn import mlx.optimizers as optim import mnist device = mx.gpu mx.set_default_device(device) mx.default_device() # Device(gpu, 0) Dataset MNIST dataset contains 60,000 training samples of 28x28 images. Thes...