Update miche/michelangelo/models/modules/checkpoint.py

miche/michelangelo/models/modules/checkpoint.py

@@ -20,10 +20,6 @@ def checkpoint(
     # :param flag: if False, disable gradient checkpointing.
     # :param use_deepspeed: if True, use deepspeed
     if flag:
-        if use_deepspeed:
-            import deepspeed
-            return deepspeed.checkpointing.checkpoint(func, *inputs)
-
         args = tuple(inputs) + tuple(params)
         return CheckpointFunction.apply(func, len(inputs), *args)
     else: