Test lm_inference_test.py fails

Hello team,

Please i need help to solve this issue, the test is failing:

python lm_inference_test.py --meliad_path=$MELIAD_PATH --data_path=$DATA

I0130 03:37:30.642391 139830854076224 nn_components.py:329] mlp: final activation = None
I0130 03:37:30.645691 139830854076224 nn_components.py:332] mlp: final = Tracedwith
I0130 03:37:30.645775 139830854076224 nn_components.py:261] mlp: residual
I0130 03:37:30.647993 139830854076224 transformer_base.py:431] tbase: pre-FFN layernorm = Tracedwith
I0130 03:37:30.648123 139830854076224 transformer_base.py:443] tbase: final FFN
I0130 03:37:30.648216 139830854076224 nn_components.py:320] mlp: hidden 4096, relu
I0130 03:37:30.651687 139830854076224 nn_components.py:329] mlp: final activation = None
I0130 03:37:30.654985 139830854076224 nn_components.py:332] mlp: final = Tracedwith
I0130 03:37:30.655071 139830854076224 nn_components.py:261] mlp: residual
I0130 03:37:30.655320 139830854076224 transformer_base.py:450] tbase: ys = Tracedwith
I0130 03:37:30.655408 139830854076224 decoder_stack.py:344] dstack: Final layernorm.
I0130 03:37:30.658026 139830854076224 decoder_stack.py:365] dstack: logits = Tracedwith
[ FAILED ] LmInferenceTest.test_lm_decode
[ RUN ] LmInferenceTest.test_lm_score_may_fail_numerically_for_external_meliad
[ FAILED ] LmInferenceTest.test_lm_score_may_fail_numerically_for_external_meliad
======================================================================
ERROR: test_lm_score_may_fail_numerically_for_external_meliad (__main__.LmInferenceTest)
----------------------------------------------------------------------
Traceback (most recent call last):
File "/home/cmgi/alphageometry/lm_inference_test.py", line 77, in test_lm_score_may_fail_numerically_for_external_meliad
outputs = LmInferenceTest.loaded_lm.beam_decode(
File "/home/cmgi/alphageometry/lm_inference.py", line 171, in beam_decode
metrics_np = self.call(inputs, dstate=dstate, eos=eos, mask=mask)
File "/home/cmgi/alphageometry/lm_inference.py", line 143, in call
_, metrics_np = self.task.run_step(self.tstate, x, 0)
File "/home/cmgi/alphageometry/meliad_lib/meliad/training_task.py", line 155, in run_step
(tstate, metrics) = self.step_function(tstate, x, self.prng_keys)
ValueError: INVALID_ARGUMENT: Invalid buffer passed: buffer has been deleted or donated.

======================================================================
FAIL: test_lm_decode (__main__.LmInferenceTest)
----------------------------------------------------------------------
Traceback (most recent call last):
File "/home/cmgi/alphageometry/lm_inference_test.py", line 71, in test_lm_decode
self.assertEqual(
AssertionError: Lists differ: [' ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇[3205 chars] ⁇ '] != ['e : D a b c e 02 D a c b e 03 ;', 'e : C [18 chars]3 ;']

First differing element 0:
' ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇ ⁇[1667 chars]\x13'
'e : D a b c e 02 D a c b e 03 ;'

Diff is 3532 characters long. Set self.maxDiff to None to see it.

----------------------------------------------------------------------
Ran 2 tests in 143.173s

FAILED (failures=1, errors=1)