Add jax_memory_cleanup argument in Task.

SamanehSaadat · Apr 12, 2024 · 677de72 · 677de72
1 parent d71f3ee
commit 677de72
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/keras_nlp/models/task.py b/keras_nlp/models/task.py
@@ -238,7 +238,7 @@ def from_preset(
             task_preset_cls = check_config_class(preset, TASK_CONFIG_FILE)
             task = load_serialized_object(preset, TASK_CONFIG_FILE)
             if load_weights:
-                jax_memory_cleanup()
+                jax_memory_cleanup(task)
                 if check_file_exists(preset, TASK_WEIGHTS_FILE):
                     task.load_task_weights(get_file(preset, TASK_WEIGHTS_FILE))
                 task.backbone.load_weights(get_file(preset, MODEL_WEIGHTS_FILE))