diff --git a/setup.py b/setup.py
index e09b714c..5e002e05 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,7 @@
         install_requires.append(line)
 
 setup(name="weaver-core",
-      version='0.4.14',
+      version='0.4.15',
       description="A streamlined deep-learning framework for high energy physics",
       long_description_content_type="text/markdown",
       author="H. Qu, C. Li",
diff --git a/weaver/train.py b/weaver/train.py
index 3dfdcb1a..e0bd0ec6 100644
--- a/weaver/train.py
+++ b/weaver/train.py
@@ -881,7 +881,7 @@ def _main(args):
             del test_loader
 
             if args.predict_output:
-                if not os.path.dirname(predict_output):
+                if not os.path.dirname(args.predict_output):
                     predict_output = os.path.join(
                         os.path.dirname(args.model_prefix),
                         'predict_output', args.predict_output)