kwywk prediction on single example looks odd

[hvgazula]

Context: Trained on one example and tested on the same example for 200 epochs (with early stopping). While the loss curves look okay, the dice curves look odd.

animations are at: /om2/scratch/Thu/for_satra/kwyk_test_2024_05_08_T0_50

Potential sources of error

1. patching (creating blocks) and stitching (putting blocks together) - check if tests exist for this.
2. predict function
3. dice function