Hi Akshay,
I’d probably look at two diagnostic (ha) questions:

If you permute the labels (labels = labels.flip(1) or so), does it still happen this way or can you then predict another organ?
If you stick with a single batch, can you overfit?

The former is a check whether there is something wrong in the tensor shapes or somesuch. I tend to print all shapes to debug that type of thing as a first step.
The second would go in the direction of capacity.
Best regards
Thomas