modified: dsm/__init__.py

	modified:   dsm/dsm_torch.py
	modified:   dsm/utilities.py

dsm/__init__.py

@@ -184,3 +184,4 @@
 from dsm.dsm_api import DeepSurvivalMachines
 from dsm.dsm_api import DeepConvolutionalSurvivalMachines
 from dsm.dsm_api import DeepRecurrentSurvivalMachines
+from dsm.dsm_api import DeepCNNRNNSurvivalMachines

dsm/dsm_torch.py

@@ -467,8 +467,8 @@ class DeepCNNRNNSurvivalMachinesTorch(DeepRecurrentSurvivalMachinesTorch):
 
   Parameters
   ----------
-  inputdim: int
-      Dimensionality of the input features.
+  inputdim: tuple
+      Dimensionality of the input features. (height, width)
   k: int
       The number of underlying parametric distributions.
   layers: int
@@ -510,16 +510,18 @@ def __init__(self, inputdim, k, typ='LSTM', layers=1,
 
     self._init_dsm_layers(hidden)
 
+    self.cnn = create_conv_representation(inputdim, hidden)
+
     if self.typ == 'LSTM':
-      self.embedding = nn.LSTM(inputdim, hidden, layers,
-                               bias=False, batch_first=True)
+      self.rnn = nn.LSTM(hidden, hidden, layers,
+                         bias=False, batch_first=True)
     if self.typ == 'RNN':
-      self.embedding = nn.RNN(inputdim, hidden, layers,
-                              bias=False, batch_first=True,
-                              nonlinearity='relu')
+      self.rnn = nn.RNN(hidden, hidden, layers,
+                        bias=False, batch_first=True,
+                        nonlinearity='relu')
     if self.typ == 'GRU':
-      self.embedding = nn.GRU(inputdim, hidden, layers,
-                              bias=False, batch_first=True)
+      self.rnn = nn.GRU(hidden, hidden, layers,
+                        bias=False, batch_first=True)
 
   def forward(self, x, risk='1'):
     """The forward function that is called when data is passed through DSM.

dsm/utilities.py

@@ -183,13 +183,14 @@ def train_dsm(model,
     costs.append(float(valid_loss))
     dics.append(deepcopy(model.state_dict()))
 
-    if (costs[-1] >= oldcost) is True:
+    if costs[-1] >= oldcost:
       if patience == 2:
         minm = np.argmin(costs)
         model.load_state_dict(dics[minm])
 
         del dics
         gc.collect()
+
         return model, i
       else:
         patience += 1
@@ -198,4 +199,10 @@ def train_dsm(model,
 
     oldcost = costs[-1]
 
+  minm = np.argmin(costs)
+  model.load_state_dict(dics[minm])
+
+  del dics
+  gc.collect()
+
   return model, i