Skip to content

Jax image classifier test

TestJaxImageClassifier #

Bases: LightningModuleTests[JaxImageClassifier]

Tests for the Jax image classification algorithm.

This simply reuses all the tests in the base test suite, specifying that the datamodule passed to the JaxImageClassifier should be for image classification and the network should be a flax.linen.Module.

test_demo #

test_demo(tmp_path: Path)

Test the demo at the bottom of the module.