Extending nn.Module class

In the video, we have extended the ‘nn.module’ class for re-shaping the tensor. But in other problems, how will we know when to extend the ‘nn.module’ class and for which problem ?

We are extending nn.Module just to reshape a tensor. We’re doing it to build a model, and reshape is part of it.

You would want to extend nn.Module whenever you wanna write a custom Neural Network model https://pytorch.org/docs/stable/generated/torch.nn.Module.html

Why have we defined a custom model?

How does working directly as below affect the output?

for images,labels in train_loader:
output = model(images.reshape(-1,784))