Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "e7696e20f9a932bad2d462d4f7cc9cf0b5000c02"
Unverified Commit ae6b134c authored by Sergii Dymchenko's avatar Sergii Dymchenko Committed by GitHub
Browse files

Pass use_reentrant for checkpoint (#8180)

parent 26fb5efe
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
torchvision/models/densenet.py
View file @ ae6b134c
...@@ -61,7 +61,7 @@ class _DenseLayer(nn.Module): ...@@ -61,7 +61,7 @@ class _DenseLayer(nn.Module):
def closure(*inputs): def closure(*inputs):
return self.bn_function(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 @torch.jit._overload_method # noqa: F811
def forward(self, input: List[Tensor]) -> Tensor: # noqa: F811 def forward(self, input: List[Tensor]) -> Tensor: # noqa: F811
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment