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))