Update app.py

app.py

@@ -18,6 +18,8 @@ import torch
 import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader, Dataset
+from sklearn.model_selection import train_test_split
+
 
 plt.switch_backend('agg')