Update lvdm/common.py

lvdm/common.py

@@ -77,7 +77,7 @@ def init_(tensor):
     tensor.uniform_(-std, std)
     return tensor
 
-ckpt = torch.utils.checkpoint.checkpoint
+# ckpt = torch.utils.checkpoint.checkpoint
 def checkpoint(func, inputs, params, flag):
     """
     Evaluate a function without caching intermediate activations, allowing for
@@ -88,7 +88,7 @@ def checkpoint(func, inputs, params, flag):
                    explicitly take as arguments.
     :param flag: if False, disable gradient checkpointing.
     """
-    if flag:
-        return ckpt(func, *inputs, use_reentrant=False)
-    else:
-        return func(*inputs)
+    # if flag:
+    #     return ckpt(func, *inputs, use_reentrant=False)
+    # else:
+    return func(*inputs)