Challenge: Distributed Training with JAX and Flax
Test your understanding of image classification and distributed training.
We will perform distributed training using JAX and Flax in this challenge. We have imported all the necessary libraries for you.
Challenge 1: Load the dataset
In the /usr/local/notebooks directory, we have a dataset in a zipped folder, cars_and_bikes.zip, containing images from two classes: cars and bikes. There are ...