Working with images, Image Shaping, Custom Model vs Resnet

I wanted to play around with a simpler model than resnet, which is what was used in the chapter 4 of the book. In the book, they create a DataLoader from the folder and use resnet18 into the learner. I used the below code and was able to get it to train. To do so, I had to reshape the data and tweak params. This breaks the ability for the dataloader show_batch. The book doesn’t have this issue, and I think it’s that resnet handles it more gracefully. What tweaks should I make to my code so the shapes line up & show_batch still works?