fix to weights path

m-lyon · Apr 11, 2022 · 98bf220 · 98bf220
1 parent 8baf320
commit 98bf220
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/dmri_rcnn/core/weights.py b/dmri_rcnn/core/weights.py
@@ -114,8 +114,8 @@ def get_weights(model_dim: int, shell: Union[int, str], q_in: int) -> str:
             Will raise a RuntimeError if not found.
     '''
     try:
-        if shell == 'all':
-            weight_dir = os.path.join(LOCAL_DIR, f'{model_dim}D_RCNN', f'{q_in}in')
+        if shell in {'all', 'all_norm'}:
+            weight_dir = os.path.join(LOCAL_DIR, f'{model_dim}D_RCNN', f'{shell}_{q_in}in')
         else:
             weight_dir = os.path.join(
                 LOCAL_DIR, f'{model_dim}D_RCNN', f'b{shell}_{q_in}in'

diff --git a/setup.py b/setup.py
@@ -12,7 +12,7 @@
     'tqdm',
 ]
 
-version = '0.3.0'
+version = '0.3.1'
 this_dir = os.path.abspath(os.path.dirname(__file__))
 with open(os.path.join(this_dir, 'README.md'), encoding='utf-8') as f:
     long_description = f.read()