Transfer learning CNNs with JAX, part 1