I basically defined a model with a Conv2D and linear layer with PyTorch and trained it with a sample dataset. The model seems to run and converge. But I am wondering I did n