Fix typing behaviour of nn.Module for mypy

[omar-abdelgawad]

Describe the bug
Currently pytorch type annotates all of its nn.Module class as returning typing.Any which isn't helpful. When you try to define your forward method and type annotate it that doesn't reflect to call of the class as mentioned in this issue (https://discuss.pytorch.org/t/adding-typing-to-call-of-nn-module/118295). Please look into the problem for how to solve it and if there is no solution maybe use the solution in the above issue.

Expected behavior
using annotating forward as def forward(self, x: torch.Tensor) -> torch.Tensor: will make the call of out = model(x) return out as a torch.Tensor