- [x] Make it - [ ] Refactor it a bit so it fits better (in terms of input/output tensors) and only uses the variables it needs.