Type fixes for type: ignore introduced in #405

[svij-sc]

W/ updating pytorch version we need to address some typecheck issues that are not addressed as part of this change

In #405, for the sake of not increasing the scope tremendously, we introduced
# type: ignore # https://github.com/Snapchat/GiGL/issues/408 where the changes to fix the type checks were not trivial.

This will need to be addressed