ensure model is in device after sparsifying

edgeml/pytorch/graph/rnn.py

@@ -231,7 +231,6 @@ class FastGRNNCell(nn.Module):
                 self.U2_old = torch.FloatTensor(np.copy(self.U2.data.cpu().detach().numpy()))
                 self.U1_old.to(self.U1.device)
                 self.U2_old.to(self.U2.device)
-
         
     def sparsify(self):
         if self._wRank is None:

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: