Set total_epochs using standard epochs keyword

Author: jmohdyusof

examples/image-vae/image_vae_baseline_pytorch.py

@@ -97,7 +97,7 @@ def run(gParams):
     test_file = candle.fetch_file(data_url + test_data, subdir='Examples/image_vae')
 
     starting_epoch = 1
-    total_epochs = None
+    total_epochs = gParams['epochs']
 
     rng_seed = 42
     torch.manual_seed(rng_seed)
@@ -263,7 +263,7 @@ def test(epoch, args):
     if total_epochs is None:
         trn_rng = itertools.count(start=starting_epoch)
     else:
-        trn_rng = range(starting_epoch, total_epochs)
+        trn_rng = range(starting_epoch, total_epochs + 1)
 
     for epoch in trn_rng:
         for param_group in optimizer.param_groups:

examples/image-vae/image_vae_default_model.txt

@@ -3,6 +3,7 @@ data_url = 'ftp://ftp.mcs.anl.gov/pub/candle/public/benchmarks/Examples/image_va
 train_data = 'train.csv'
 test_data = 'test.csv'
 workers = 16
+epochs = None
 batch_size = 256
 grad_clip = 2.0
 model_path = 'models'