add deepspeed to try-except block

ttengwang · Jun 11, 2023 · 64ce3e0 · 64ce3e0
1 parent a48d27c
commit 64ce3e0
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/EVA-CLIP/rei/eva_clip/eva_vit_model.py b/EVA-CLIP/rei/eva_clip/eva_vit_model.py
@@ -16,7 16,10 @@
 from .rope import VisionRotaryEmbedding, VisionRotaryEmbeddingFast
 
 if os.getenv('ENV_TYPE') == 'deepspeed':
- from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
  try:
  from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
  except:
  from torch.utils.checkpoint import checkpoint
 else:
  from torch.utils.checkpoint import checkpoint