Update main.py

MadryLab · Aug 21, 2019 · 57f872a · 57f872a
1 parent 00b1f27
commit 57f872a
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/robustness/main.py b/robustness/main.py
@@ -17,7 +17,7 @@
     from .datasets import DATASETS
     from .train import train_model, eval_model
     from .tools import constants, helpers
-    from . import defaults
+    from . import defaults, __version__
     from .defaults import check_and_fill_args
 except:
     raise ValueError("Make sure to run with python -m (see README.md)")
@@ -90,11 +90,10 @@ def setup_store_with_metadata(args):
     try:
         repo = git.Repo(path=os.path.dirname(os.path.realpath(__file__)),
                             search_parent_directories=True)
-        git_commit = repo.head.object.hexsha
+        version = repo.head.object.hexsha
     except git.exc.InvalidGitRepositoryError:
-        git_commit = 'pip_' + helpers.get_pip_version()
-        print(git_commit)
-    args.git_commit = git_commit
+        version = __version__
+    args.version = version
 
     # Create the store
     store = cox.store.Store(args.out_dir, args.exp_name)