GitLab

* Start adding types

* add typing

* Type prototype models

* fix optional type bug

* transient import

* Fix weights type

* fix import

* Apply suggestions from code review

Address nits
Co-authored-by: Vasilis Vryniotis <datumbox@users.noreply.github.com>
Co-authored-by: Khushi Agrawal <khushiagrawal411@gmail.com>
Co-authored-by: Vasilis Vryniotis <datumbox@users.noreply.github.com>