Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit ae6b134

Browse files
authoredDec 22, 2023
Pass use_reentrant for checkpoint (#8180)
1 parent 26fb5ef commit ae6b134

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed
 

‎torchvision/models/densenet.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ def call_checkpoint_bottleneck(self, input: List[Tensor]) -> Tensor:
6161
def closure(*inputs):
6262
return self.bn_function(inputs)
6363

64-
return cp.checkpoint(closure, *input)
64+
return cp.checkpoint(closure, *input, use_reentrant=False)
6565

6666
@torch.jit._overload_method # noqa: F811
6767
def forward(self, input: List[Tensor]) -> Tensor: # noqa: F811

0 commit comments

Comments
 (0)
Please sign in to comment.