text: Title: [Tracker] Split tests into multiple file tests
  Body: Currently, we have some test files that have all (or almost all) tests of a module within a single file, the ideal here is to split these cases into multiple test files. Which file should be testing a single functionality, or file of our library.
  
  Example:
  > We have now a single file testing almost all losses functions https://github.com/kornia/kornia/blob/c5558ac2c03bbbbbbb6c9f82616b8d5229398f1a/tests/losses/test_losses.py we would like to have a single test file for each loss. What would mirror the file organization of the losses module of kornia https://github.com/kornia/kornia/tree/c5558ac2c03bbbbbbb6c9f82616b8d5229398f1a/kornia/losses
  
  This issue will work as a tracker to map some of these files that need to be split, feel free to help us find others and/or help us split them into multiple files.
  
  The goal here is to reorganize our test suite. If you have any questions, you can call us and/or ask on our Slack page too [![Slack](https://img.shields.io/badge/Slack-4A154B?logo=slack&logoColor=white)](https://join.slack.com/t/kornia/shared_invite/zt-csobk21g-2AQRi~X9Uu6PLMuUZdvfjA)
  
  _Originally posted by @edgarriba in https://github.com/kornia/kornia/pull/2745#discussion_r1461605306_
  
  -------------
  You can choose anyone below, or find new ones. You don't have to worry about separating all the tests in a file into a single PR
  
  - [x] Split losses tests - https://github.com/kornia/kornia/blob/ce434e467faf617604bb3383cf78cd0b79f59dbd/tests/losses/test_losses.py - #2801
  - [x] Split metrics tests - https://github.com/kornia/kornia/blob/ce434e467faf617604bb3383cf78cd0b79f59dbd/tests/test_metrics.py
  - [x] Split contrib tests - https://github.com/kornia/kornia/blob/ce434e467faf617604bb3383cf78cd0b79f59dbd/tests/contrib/test_contrib.py  - #2802
  - [ ] Split augmentations tests
    - [ ] 2d - https://github.com/kornia/kornia/blob/ce434e467faf617604bb3383cf78cd0b79f59dbd/tests/augmentation/test_augmentation.py
    - [ ] 3d - https://github.com/kornia/kornia/blob/ce434e467faf617604bb3383cf78cd0b79f59dbd/tests/augmentation/test_augmentation_3d.py
  - [ ] Split container augmentation tests - https://github.com/kornia/kornia/blob/ce434e467faf617604bb3383cf78cd0b79f59dbd/tests/augmentation/test_container.py