Add tensor typing

[juanmc2005]

Problem

It's difficult to understand the shape and format expected from tensors.

Idea

Add torchtyping. There's also the possibility to check that types are respected at runtime (instead of it being for documentation only), for example while debugging.