Can't increase bs for Vista3d bundle finetuning

[AHarouni]

Describe the bug
I am using vista3d fine tunning for a new task where I have 3 labels. It runs fine with bs=1 and it uses 8GB. I have a 48 GB gpu so I tried increasing the BS but I get an error. I also tried increasing the number of patched per image and got the same error

To Reproduce
download vista3d bundle

1. change number of classes to 3
2. change dataset json to your data as liver tumor dataset
3. run fine tuning, it would run fine
4. now change the bs to 2, rerun you will get an error

Expected behavior
I should be able to increase the BS >1 and make use of my GPU memory

Additional context
the GPU utilization is great it is always 90%+