check if we can implement torch.Module

[shashvatshah9]

extend and overload https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module