TorchModuleWrapper serialization issue

[MicheleCattaneo]

I would like to open this issue to revive a discussion started in a previous issue (#19226). While the previous issue seems to be inactive, the potential bug seems to still be present. I hope this is fine.

The problem arises when trying to save a mdel containing a TorchModuleWrapper layer (therefore using PyTorch as backend).

I referenced the original issue and in particular my latest comment below for more details:

This bug is currently still present. The following is a minimal snippet that can reproduce it:

import os
os.environ["KERAS_BACKEND"] = "torch"
import torch
import keras

torch_module = torch.nn.Linear(4,4)
keras_layer = keras.layers.TorchModuleWrapper(torch_module)

inputs = keras.Input(shape=(4,))
outputs = keras_layer(inputs)
model = keras.Model(inputs=inputs, outputs=outputs)

model.save('./serialized.keras')

The error is:

UnicodeDecodeError: 'utf-8' codec can't decode byte 0x80 in position 64: invalid start byte

generated in keras.src.saving.serialization_lib.serialize_keras_object

It is worth noting that manually using get_config and from_config to serialize and deserialize (in memory) produce the correct result:

torch_linear = torch.nn.Linear(4,4) # xA^T+b with initalized weights
wrapped_torch = TorchModuleWrapper(torch_linear) # Wrap it 

# get its config, and rebuild it 
torch_linear_from_config = keras.layers.TorchModuleWrapper.from_config(wrapped_torch.get_config()).module

# assert all parameters are the same
assert (torch_linear.weight == torch_linear_from_config.weight).all()
assert (torch_linear.bias == torch_linear_from_config.bias).all()

What get_config() does is map module (a torch object) to its serialized string (coming from torch.save(self.module, buffer)). I believe it is wrong to use the utf-8 in serialize_keras_object(obj), since that encoding is specifically meant for text and not arbitrary bytes.

Does anybody have an idea about it?
Thank you for any help on this!

I got this error with both:

• python 3.10, keras 3.7.0, torch 2.5.1+cu124
• python 3.11, keras 3.8.0, torch 2.5.1+cu124

Originally posted by @MicheleCattaneo in #19226

As I am highly interested in using Keras3 with PyTorch modules, I am willing to contribute to a potential solution to this issue. I would however appreciate some guidance, as I am not very familiar with the Keras code base.

Thank you for any help!