@shravan I took another pass at your code and figured out some errors but trying to correct one led to another and so on and I couldn’t take it further in the interest of time (sorry about that). But, I’ll try to help you start so you know what exactly is the issue. One, the input image which is 224x224x3 is transformed using ToPILImage and further flattened using view. Try to understand the shape of the output (3x224x3 = 2016). Later on, when you define your EmergencyVehicleDetectionModel class you hard-coded an input of size 224x224x3 = 150528. I hope you can see what’s happening here. You are feeding a tensor of shape 2016 to a network expecting a tensor of shape 150528. I hope this will help you understand where to look in your code and how to go about it. The other follow-up issues are not really difficult to address and some intuitive understanding of how various NN layers and also loss functions work and what shapes are expected will come in very handy. I urge you to spend more time going through PyTorch documentation for every function you use. I wish you the best.