HI Miguel,
I have just tried the modified codes and no error was raised.
From model.named_modules() ,I can see the difference of these two networks:
As for original Net:
(conv1): Conv2d(1, 10, kernel_size=(5, 5), stride=(1, 1))
(conv2): Conv2d(10, 20, kernel_size=(5, 5), stride=(1, 1))
(conv2_drop): Dropout2d (p=0.5)
(fc1): Linear (320 -> 50)
(fc2): Linear (50 -> 10)

As for the modified version:
('conv1', Conv2d(1, 10, kernel_size=(5, 5), stride=(1, 1)))
('conv2', Conv2d(10, 20, kernel_size=(5, 5), stride=(1, 1)))
('fc1', Linear (320 -> 50))
('fc2', Linear (50 -> 10))

There is a dropout layer missed. But both network models works and should behave the same operation in forward propagation.
I wonder if these two networks are naturally the same during training and testing?
@apaszke @smth @colesbury
Hope you can also give me some suggestions if possible.
Thanks.