Faster batch scaler tests

[SkafteNicki]

As discussed on slack the two slowest test has to do with the batch scaler.
This PR tries to bring their test time down.

[SkafteNicki]

@Borda I moved some of the batch scaling test from cpu to gpu (where they probably should have been all the time). It reduces the time of the two longest test by around 40 sec each, so total reduction by 1:30 min.

[Borda]

@SkafteNicki can we also make the test smaller, just less data or larger steps?

[SkafteNicki]

we could limit the size of the training set since after PR #3271 the search will stop when batch_size>=len(train_dataset), meaning lower test time

[awaelchli]

@SkafteNicki I was thinking the same yesterday but then on the other hand, part of the test is also to check that the OOM handling goes correctly. so we would sacrifie that.

[SkafteNicki]

@awaelchli also what I was thinking. Only other way to make the search stop earlier if we artificially filled up the gpu memory before running the search. Else let keep it as it is now (still reduced significantly compared to before).