Pass use_reentrant for checkpoint (#8180)

pytorch · Dec 22, 2023 · ae6b134 · ae6b134
1 parent 26fb5ef
commit ae6b134
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchvision/models/densenet.py b/torchvision/models/densenet.py
@@ -61,7 +61,7 @@ def call_checkpoint_bottleneck(self, input: List[Tensor]) -> Tensor:
         def closure(*inputs):
             return self.bn_function(inputs)
 
-        return cp.checkpoint(closure, *input)
+        return cp.checkpoint(closure, *input, use_reentrant=False)
 
     @torch.jit._overload_method  # noqa: F811
     def forward(self, input: List[Tensor]) -> Tensor:  # noqa: F811