reduce time on neurochem trainer on CI (#208)

tests/test_data/inputtrain.ipt

@@ -2,7 +2,7 @@
 sflparamsfile=rHCNO-5.2R_16-3.5A_a4-8.params
 ntwkStoreDir=networks/
 atomEnergyFile=sae_linfit.dat
-nmax=0! Maximum number of iterations (0 = inf)
+nmax=10! Maximum number of iterations (0 = inf)
 tolr=1! Tolerance - early stopping
 emult=0.1!Multiplier by eta after tol switch
 eta=0.001! Eta -- Learning rate

tests/test_data/inputtrain.yaml

@@ -2,7 +2,7 @@
 sflparamsfile: 'rHCNO-5.2R_16-3.5A_a4-8.params'
 ntwkStoreDir: 'networks/'
 atomEnergyFile: 'sae_linfit.dat'
-nmax: 0 # Maximum number of iterations (0 = inf)
+nmax: 10 # Maximum number of iterations (0 = inf)
 tolr: 1 # Tolerance - early stopping
 emult: 0.1 # Multiplier by eta after tol switch
 eta: 0.001 # Eta -- Learning rate

torchani/neurochem/__init__.py

@@ -556,7 +556,6 @@ if sys.version_info[0] > 2:
             assert_param('runtype', 'ANNP_CREATE_HDNN_AND_TRAIN')
             assert_param('adptlrn', 'OFF')
             assert_param('tmax', 0)
-            assert_param('nmax', 0)
             assert_param('ntwshr', 0)
 
             # load parameters
@@ -585,6 +584,8 @@ if sys.version_info[0] > 2:
             del params['tbtchsz']
             self.validation_batch_size = params['vbtchsz']
             del params['vbtchsz']
+            self.nmax = params['nmax']
+            del params['nmax']
 
             # construct networks
             input_size, network_setup = network_setup
@@ -700,6 +701,12 @@ if sys.version_info[0] > 2:
                     self.global_epoch = trainer.state.epoch
                     self.global_iteration = trainer.state.iteration
 
+                if self.nmax > 0:
+                    @trainer.on(ignite.engine.Events.EPOCH_COMPLETED)
+                    def terminate_when_nmax_reaches(trainer):
+                        if trainer.state.epoch >= self.nmax:
+                            trainer.terminate()
+
                 if self.tqdm is not None:
                     @trainer.on(ignite.engine.Events.EPOCH_STARTED)
                     def init_tqdm(trainer):