Challenge: Train and Test a Model

Learn to implement a full training and testing pipeline for a CNN model in JAX. Explore dataset loading, loss and metric definitions, training state management, and evaluation steps to solidify your practical skills.