Colby Banbury

ValueError: Received incompatible devices for jitted computation. (demo_jax_distributed.py)

2

I get the following error when running `examples/demo_jax_distributed.py` on a Cloud TPU VM (tpuv2): ``` Traceback (most recent call last): File "/home/colby/jax_test.py", line 341, in loss, accuracy = model.evaluate(eval_data) File...

The training script defines both BATCH_SIZE and batch_size but only uses BATCH_SIZE. I suggest we remove the unused batch_size to avoid confusion. https://github.com/mlcommons/tiny/blob/3bca1a3adc7d9ec5186c0ff376715f2d3b19f382/benchmark/training/visual_wake_words/train_vww.py#L19 https://github.com/mlcommons/tiny/blob/3bca1a3adc7d9ec5186c0ff376715f2d3b19f382/benchmark/training/visual_wake_words/train_vww.py#L33

adds custom topic selection

1

This is definitely a hack, but it works