Your code seems to run fine on my machine using random inputs of 1x28x28.
Could you post the error message and in which line of code it is thrown?
PS: you can add code snippets using three backticks ``` 
I’ve formatted your code, so that I could easily copy it.