Getting error when trying to export GemmaCausalLM model to SavedModel format

[freshbash]

Hi!

I'm working on a project where I train the Gemma 2B model. The training part went smoothly. The problem comes when I try to export the model to SavedModel format so that I can deploy it to GCP for demo purposes. I'm using the following code to load and export the model:

model = keras_nlp.models.GemmaCausalLM.from_preset("kaggle://<username>/<model>")
model.export("saved_model")

When the model is loaded on system RAM and I try to export the it, I'm getting the following error:

ValueError                                Traceback (most recent call last)
[<ipython-input-15-e773bda33c50>](https://localhost:8080/#) in <cell line: 1>()
----> 1 model.export("saved_model")

27 frames
    [... skipping hidden 10 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/lax_numpy.py](https://localhost:8080/#) in _arange_dynamic(start, stop, step, dtype)
   6567   # Here if at least one of start, stop, step are dynamic.
   6568   if any(not core.is_dim(v) for v in (start, stop, step)):
-> 6569     raise ValueError(
   6570         "In arange with non-constant arguments all of start, stop, and step "
   6571         f"must be either dimension expressions or integers: start={start}, "

ValueError: Exception encountered when calling GemmaDecoderBlock.call().

In arange with non-constant arguments all of start, stop, and step must be either dimension expressions or integers: start=b, stop=None, step=1

Arguments received by GemmaDecoderBlock.call():
  • x=jnp.ndarray(shape=(batch, b, 2048), dtype=float32)
  • padding_mask=jnp.ndarray(shape=(batch, a), dtype=float32)
  • cache=None
  • cache_update_index=0

And when the model is loaded on GPU, I get this error:

InternalError: Failed copying input tensor from /job:localhost/replica:0/task:0/device:CPU:0 to /job:localhost/replica:0/task:0/device:GPU:0 in order to run _EagerConst: Dst tensor is not initialized.

Now, I'm not really sure why is this happening because everything before this step(training, inference) worked perfectly. I'm also new to keras and DL frameworks in general. Any sort of advice would be really helpful.

Regards,
Bhashwar Sengupta