Is your feature request related to a problem? Please describe.

Currently, the engines (in engines.py) are not torch.nn.modules and not all of them implement a forward function. This complicates the code in model.py and leads to other unexpected issues.

Describe the solution you'd like

See above. Any apparent inconsistencies should also be fixed there.

Describe alternatives you've considered

No response

Additional context

No response

Organisation

No response