[FEATURE] speed up CI tests

[VukW]

Is your feature request related to a problem? Please describe.

Now running all the tests takes around ~30 mins. I understand we run a few training loops (on CPU) during tests; however, maybe it is still possible to speed them up?

Describe the solution you'd like

• Profile tests:
  • Are there any specific tests that takes bigger part of time or do all tests have relatively the same amount of time?
  • Profile slow tests: is it just a training loop that takes time? or do we make any additional transformations or something else?
  • Tune slow tests: reduce dataset size? reduce number of epochs? something else?

Crucial is to keep the same tests coverage while tuning tests.