Update `train_model` to automatically compute NCLASSES from data

[dbuscombe-usgs]

See: Doodleverse/doodleverse_utils#10 (comment)

New dice and IoU functions are agnostic to, and do not check, correct specification of NCLASSES. Additionally, cat and kld loss (and presumably others) fail when NCLASSES is incorrectly specified due to operator error

The above thread proposes a simple solution, that has been tested and works. Should appear in the next release.

[dbuscombe-usgs]

The code block should be inserted after

val_ds = val_ds.prefetch(AUTO) #

[dbuscombe-usgs]

this is already implemented in this branch of the code: https://github.com/Doodleverse/segmentation_gym/tree/new_makedatasets