Tests

Update notebook slightly #283

Job	Run time
lint	1m 17s
test (ubuntu-latest, 3.10, jax)	5m 47s
test (ubuntu-latest, 3.10, tensorflow)	5m 42s
test (ubuntu-latest, 3.10, torch)	4m 13s
test (ubuntu-latest, 3.11, jax)	6m 46s
test (ubuntu-latest, 3.11, tensorflow)	5m 44s
test (ubuntu-latest, 3.11, torch)	4m 7s
test (windows-latest, 3.10, jax)	12m 5s
test (windows-latest, 3.10, tensorflow)	10m 54s
test (windows-latest, 3.10, torch)	10m 52s
test (windows-latest, 3.11, jax)	12m 8s
test (windows-latest, 3.11, tensorflow)	13m 17s
test (windows-latest, 3.11, torch)	11m 19s
	1h 44m 11s

Provide feedback