Update train_mlp.py

train_mlp.py

@@ -38,9 +38,9 @@ def train_model(model, train_dataset, val_dataset, epochs=10, lr=0.001, save_los
         running_loss = 0.0
         for example in train_dataset:
             img = example['image']
-            img = np.array(img)  # Convert PIL image to NumPy array
-            img = img.transpose((2, 0, 1))  # Transpose to (channels, height, width)
-            img = torch.from_numpy(img).float().reshape(1, -1).to(device)  # Convert to tensor and reshape
+            img = np.array(img.convert('L'))  # Convert PIL image to grayscale NumPy array
+            img = img.reshape(1, -1)  # Flatten the image
+            img = torch.from_numpy(img).float().to(device)  # Convert to tensor
             label = torch.tensor([example['label']]).to(device)
 
             optimizer.zero_grad()
@@ -63,9 +63,9 @@ def train_model(model, train_dataset, val_dataset, epochs=10, lr=0.001, save_los
         with torch.no_grad():
             for example in val_dataset:
                 img = example['image']
-                img = np.array(img)  # Convert PIL image to NumPy array
-                img = img.transpose((2, 0, 1))  # Transpose to (channels, height, width)
-                img = torch.from_numpy(img).float().reshape(1, -1).to(device)  # Convert to tensor and reshape
+                img = np.array(img.convert('L'))  # Convert PIL image to grayscale NumPy array
+                img = img.reshape(1, -1)  # Flatten the image
+                img = torch.from_numpy(img).float().to(device)  # Convert to tensor
                 label = torch.tensor([example['label']]).to(device)
 
                 outputs = model(img)
@@ -99,7 +99,7 @@ def main():
 
     # Split the dataset into train and validation sets
     train_dataset = dataset['train']
-    val_dataset = dataset['valid']
+    val_dataset = dataset['validation']  # Assuming 'validation' is the correct key
 
     # Determine the number of classes
     num_classes = len(set(train_dataset['label']))
@@ -108,7 +108,7 @@ def main():
     image_size = 64  # Assuming the images are square
 
     # Define the model
-    input_size = image_size * image_size * 3
+    input_size = image_size * image_size  # Since images are grayscale
     hidden_sizes = [args.width] * args.layer_count
     output_size = num_classes