Hi @Sally_Goldin, welcome to ML with Pytorch,
Before checking the below suggestion, please make sure that you are indeed overfitting. If yes, skip the step A.
Step A. Make your Model overfit

According to my experiences, I tend to do these things:

Try to make the model overfit (which is currently not optimized on your model)
Use pretrained model (this will HUGE-ly increase the performance and the speed of your training time. You could use these available models 2 (*thanks to the community)
*Try to avoid the regularizations as much as you can at first, aka Dropout, L2/L1 norm, or other fancy ones on the internet.
In my experience, try to avoid using Batchnorm on the last layer.
Double check your final activation layer (in your case, you should have used “softmax” CE since your goal is image recognition/discrimination).
Try to add batchnorm2d on your convolution layers (in case you are using custom model).
Try to normalize your image (in case you haven’t normalized it, you could normalized it [0,1] or [-1,1] etc).
(These 7 steps should basically solved most of your initial issue, which is to make your model very good on your training data)

Step B. Make your Model generalized on the unseen data

This is where you should do generalizations*. There are many resources on the internet for your starting point to save your time.
Double check your training data’s distribution, make sure that they have similar total images per-class (this is often happening).
If it is still not working, you might gonna need to check your image, make sure that they are recognizable (this is often happening in my prev. projects).

Hope it helps, cheers~