How to define model

User-defined models can be anything derived from torch.nn.Module. For example, we can define a convolutional neural network (CNN) as follows:

In the code, one can create the CNN model and the loss function instances as follows:

model = CNN()
loss_fn = torch.nn.CrossEntropyLoss()

Note that the loss_fn can be anything derived from torch.nn.Module as well.