Update main.py

main.py

@@ -18,16 +18,17 @@ def to_numpy(examples):
     return examples
 
 def preprocess():
-    test_dataset = load_dataset("active-learning/test_mnist")
-    train_dataset = load_dataset("active-learning/labeled_samples")
+    train_dataset = load_dataset("active-learning/labeled_samples")["train"]
     train_dataset = train_dataset.map(to_numpy, batched=True)
+
+    test_dataset = load_dataset("active-learning/test_mnist")["test"]
     test_dataset = test_dataset.map(to_numpy, batched=True)
     
-    x_train = train_dataset["train"]["pixel_values"]
-    y_train = train_dataset["train"]["label"]
+    x_train = train_dataset["pixel_values"]
+    y_train = train_dataset["label"]
     
-    x_test = test_dataset["test"]["pixel_values"]
-    y_test = test_dataset["test"]["label"]
+    x_test = test_dataset["pixel_values"]
+    y_test = test_dataset["label"]
 
     x_train = np.expand_dims(x_train, -1)
     x_test = np.expand_dims(x_test, -1)