ensure model is in device after sparsifying

edgeml/pytorch/graph/rnn.py

@@ -232,7 +232,6 @@ class FastGRNNCell(nn.Module):
                 self.U1_old.to(self.U1.device)
                 self.U2_old.to(self.U2.device)
         
-        
     def sparsify(self):
         if self._wRank is None:
             self.W.data = utils.hardThreshold(self.W, self._wSparsity)

examples/pytorch/FastCells/train_classifier.py

@@ -239,6 +239,7 @@ class KeywordSpotter(nn.Module):
                                 self.sparsifyWithSupport()
                         else:
                             self.sparsifyWithSupport()
+                    self.to(device) # sparsify routines might move param matrices to cpu
 
                 learning_rate = optimizer.param_groups[0]['lr']
                 if detail: