Update app.py

app.py

@@ -274,7 +274,7 @@ class LSTMAutoencoder(nn.Module):
         return out
 
 
-def lstm_anomaly_detection(X, feature_columns, raw_embedding_columns, epochs=100):
+def lstm_anomaly_detection(X, feature_columns, raw_embedding_columns, epochs=100, batch_size=64):
     device = 'cuda'
     X = torch.FloatTensor(X).to(device)
     if X.dim() == 2:
@@ -319,8 +319,8 @@ def lstm_anomaly_detection(X, feature_columns, raw_embedding_columns, epochs=100
 
     return mse_all, mse_comp, mse_raw
 
-def embedding_anomaly_detection(embeddings, epochs=100):
-    device = 'cpu'
+def embedding_anomaly_detection(embeddings, epochs=100, batch_size=64):
+    device = 'cuda'
     X = torch.FloatTensor(embeddings).to(device)
     if X.dim() == 2:
         X = X.unsqueeze(0)