temp fix for Deepspeed related issue

Azure · Jul 26, 2024 · 3de2c08 · 3de2c08
1 parent e7b23f4
commit 3de2c08
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/assets/training/finetune_acft_image/src/finetune/finetune.py b/assets/training/finetune_acft_image/src/finetune/finetune.py
@@ -3,6 +3,9 @@
 
 """File containing function for finetune component."""
 
+# Remove when base image issues are resolved.
+import deepspeed
+
 import os
 import json
 import argparse
@@ -881,6 +884,8 @@ def main():
     parser = get_parser()
     args, _ = parser.parse_known_args()
 
+    print(f"Deepspeed Version: {deepspeed.__version__}")
+
     if args.task_name in [Tasks.HF_SD_TEXT_TO_IMAGE]:
         parser = add_sd_args_to_parser(parser)
         args, _ = parser.parse_known_args()