Fix flaky test_classification_model[cuda-vit_h_14] test #5254
Description
🐛 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