Skip to content

Commit

Permalink
[gradient_checkpointing] default to use it for torch 2.3 (#28538)
Browse files Browse the repository at this point in the history
* default to use it

* style
  • Loading branch information
ArthurZucker authored Feb 20, 2024
1 parent 49c0b29 commit 9094abe
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion src/transformers/modeling_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2092,7 +2092,7 @@ def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")

if gradient_checkpointing_kwargs is None:
gradient_checkpointing_kwargs = {}
gradient_checkpointing_kwargs = {"use_reentrant": True}

gradient_checkpointing_func = functools.partial(checkpoint, **gradient_checkpointing_kwargs)

Expand Down

0 comments on commit 9094abe

Please sign in to comment.