enable use_checkpoint flag for Attention Block (#153)

atong01 · Dec 3, 2024 · bb20577 · bb20577
1 parent 581a7d9
commit bb20577
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchcfm/models/unet/unet.py b/torchcfm/models/unet/unet.py
@@ -270,7 +270,7 @@ def __init__(
         self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
 
     def forward(self, x):
-        return checkpoint(self._forward, (x,), self.parameters(), True)
+        return checkpoint(self._forward, (x,), self.parameters(), self.use_checkpoint)
 
     def _forward(self, x):
         b, c, *spatial = x.shape