tests: unit tests do not check the outputs of models

While the unit tests do cover the output types, they do not check if the outputs themselves are correct/incorrect.

To do:
* Use some loss function w.r.t ideal outputs and make sure the losses are low enough
* Make sure the output `dtype` matches the original (`int, float32, float64`)