Fix training issue and I improved the code

[andreabac3]

Hi @rslim087a thank you for your works.

I have played with this notebook, I found some issue and code portion to improve.

1. During the training loop you have missed to set the model in eval and train mode.
2. To improve the readability I moved from the x.view(-1, 4464) to the more modern pytorch function called torch.flatten(x, start_dim=1)
3. I also added a constant variable section.

future improments: add the pytorch reproducibility stuff.

Kind Regards,
andreabac3