Hey,
found the issue. the train.py expected the hyperparamter to be train-batch-size and we always passed in train_batch_size so it used the default of 32 and not the values passed in. I fixed all of the examples and tested 256 and it failed.
If you want to test fp16 you would need to change the train.py