Wesley_Neill:

nn.Linear(layers[i], layers[i + 1])


Now that makes sense. I was thinking how can one create a network combined of Conv1d and Conv2d and few other layers using just tuple of numbers of in out channels.
Let me disagree with you on this topic. Actually, one of the most fascinating things about PyTorch is its modularity and in this case, you are building a customized model, a network which will be represented as a computational graph. This model can be used for different tasks with different purposes. So, embedding loss and optimizer within train method makes it less modular. For instance if I want to use your implementation for my own data, I had to change source code directly rather than passing the final model and its outputs to a customized loss or optimizer.
If you go through more in PyTorch’s source code, you will enjoy its modularity!



 Wesley_Neill:

It is less flexible this way, but works now


If you are interested in creating high level classes and wrapping all those steps into single class, I think reading source code of PyTorch Lightening and Ignite would be a good idea.