Update train_mlp.py

train_mlp.py

@@ -24,6 +24,9 @@ class MLP(nn.Module):
 
 # Train the model
 def train_model(model, train_dataset, val_dataset, epochs=10, lr=0.001, save_loss_path=None):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=lr)
 
@@ -35,8 +38,8 @@ 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 = np.array(example['image'])
-            img = torch.from_numpy(img).float().view(1, -1)
-            label = torch.tensor([example['label']])
+            img = torch.from_numpy(img).float().view(1, -1).to(device)
+            label = torch.tensor([example['label']]).to(device)
 
             optimizer.zero_grad()
             outputs = model(img)
@@ -58,8 +61,8 @@ 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 = np.array(example['image'])
-                img = torch.from_numpy(img).float().view(1, -1)
-                label = torch.tensor([example['label']])
+                img = torch.from_numpy(img).float().view(1, -1).to(device)
+                label = torch.tensor([example['label']]).to(device)
 
                 outputs = model(img)
                 loss = criterion(outputs, label)