Update train_mlp.py

train_mlp.py

@@ -37,8 +37,10 @@ def train_model(model, train_dataset, val_dataset, epochs=10, lr=0.001, save_los
         model.train()
         running_loss = 0.0
         for example in train_dataset:
-            img = np.array(example['image'])
-            img = torch.from_numpy(img).float().view(1, -1).to(device)
+            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
             label = torch.tensor([example['label']]).to(device)
 
             optimizer.zero_grad()
@@ -60,8 +62,10 @@ def train_model(model, train_dataset, val_dataset, epochs=10, lr=0.001, save_los
         total = 0
         with torch.no_grad():
             for example in val_dataset:
-                img = np.array(example['image'])
-                img = torch.from_numpy(img).float().view(1, -1).to(device)
+                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
                 label = torch.tensor([example['label']]).to(device)
 
                 outputs = model(img)