Update generator.py

globus-labs · Oct 18, 2023 · f90d955 · f90d955
1 parent f2d1936
commit f90d955
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/mofa/generator.py b/mofa/generator.py
@@ -7,6 +7,7 @@
 from difflinker_sample import sample_from_sdf
 from difflinker_train import get_args, main
 import yaml
+import os
 
 def train_generator(
         # starting_model: str | Path,
@@ -22,6 +23,8 @@ def train_generator(
     Returns:
         Path to the new model weights
     """
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
     args = get_args()
 
     if args.config:
@@ -70,4 +73,6 @@ def run_generator(
     print("Saved XYZ files in mofa/output directory!")
 
 if __name__ == "__main__":
-    run_generator()
+    # run_generator()
+    train_generator()
+