Fix flaky `test_classification_model[cuda-vit_h_14]` test

[datumbox]

🐛 Describe the bug

After the addition of vit_h_14 (#5210) we started receiving memory issues in some of our tests. Unfortunately the don't fail consistently but rather occasionally. Examples 1 and 2.

Sometimes the memory issues are limited to that specific test, while other (see example 2) it takes down with it all other tests. We need to address this and restore our CI to green.

Versions

Latest main branch - 4bf6c6e