* Fix model loading collate_fn non-static * Fix serialization * Add tests to reproduce the error * Comment the change to let the tests fail * Finish changes