We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent de2b6bd commit 66edcb5Copy full SHA for 66edcb5
monai/networks/nets/unet.py
@@ -25,6 +25,7 @@
25
26
__all__ = ["UNet", "Unet"]
27
28
+
29
class _ActivationCheckpointWrapper(nn.Module):
30
def __init__(self, module: nn.Module) -> None:
31
super().__init__()
@@ -35,6 +36,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
35
36
return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
37
return cast(torch.Tensor, self.module(x))
38
39
40
class UNet(nn.Module):
41
"""
42
Enhanced version of UNet which has residual units implemented with the ResidualUnit class.
0 commit comments