added condition for hf to train script

boettiger-lab · Feb 28, 2024 · f00b75c · f00b75c
1 parent aefff0c
commit f00b75c
Showing 1 changed file with 11 additions and 7 deletions.
diff --git a/scripts/train.py b/scripts/train.py
@@ -6,13 +6,17 @@
 
 import rl4caribou
 
-
+# training
+#
 from rl4caribou.utils import sb3_train    
 model_save_id, train_options = sb3_train(args.file)
 
-from rl4caribou.utils import upload_to_hf
-try:
-    upload_to_hf(args.file, "sb3/"+args.file, repo=train_options['repo'])
-    upload_to_hf(model_save_id, "sb3/"+model_save_id+".zip", repo=train_options['repo'])
-except:
-    print("Couldn't upload to hf!")
+# hugging face
+#
+if 'repo' in train_options:
+    from rl4caribou.utils import upload_to_hf
+    try:
+        upload_to_hf(args.file, "sb3/"+args.file, repo=train_options['repo'])
+        upload_to_hf(model_save_id, "sb3/"+model_save_id+".zip", repo=train_options['repo'])
+    except:
+        print("Couldn't upload to hf!")