Update models/LSTM.py

models/LSTM.py

@@ -7,6 +7,7 @@ from models.preprocess_stage.preprocess_lstm import preprocess_lstm
 EMBEDDING_DIM = 128
 HIDDEN_SIZE = 16
 MAX_LEN = 125
+DEVICE='cpu'
 
 embedding_matrix = np.load('models/datasets/embedding_matrix.npy')
 embedding_layer = nn.Embedding.from_pretrained(torch.FloatTensor(embedding_matrix))
@@ -79,6 +80,7 @@ def predict_3(text):
 
     preprocessed_text = preprocess_lstm(text, MAX_LEN=MAX_LEN)
 
+    model.to(DEVICE)
     model.eval()
     predict, attention = model(torch.tensor(preprocessed_text).unsqueeze(0))