Check probe batch size 512 vs 1024

[robertdvdk]

AIM uses a batch size of 1024 for finetuning their probe, we use 512. We should check whether 1024 improves performance.