Compatibility with newer tensorflow versions

We should make the code compatible with newer tensorflow versions. This means

1. [ ] Custom layers should implement `get_config` and `from_config` methods
2. [ ] Pass `custom_object` to load_model function (is it necessary?)