[Testing] Add unit tests for data_loader.py

[TechieSamosa]

Use pytest to write a test that verifies the getitem method returns the correct (C, H, W) tensor shape.

[TechieSamosa]

If you have questions about this task, please ask in the GitHub Discussions board. Please comment 'I would like to work on this' to be assigned!