Enable GPU support; move model and data to GPU if available

Author: guojing0

beginner_source/basics/optimization_tutorial.py

@@ -64,7 +64,10 @@ def forward(self, x):
         logits = self.linear_relu_stack(x)
         return logits
 
-model = NeuralNetwork()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"Using device: {device}")
+
+model = NeuralNetwork().to(device)
 
 
 ##############################################
@@ -153,6 +156,7 @@ def train_loop(dataloader, model, loss_fn, optimizer):
     # Unnecessary in this situation but added for best practices
     model.train()
     for batch, (X, y) in enumerate(dataloader):
+        X, y = X.to(device), y.to(device)
         # Compute prediction and loss
         pred = model(X)
         loss = loss_fn(pred, y)
@@ -163,7 +167,7 @@ def train_loop(dataloader, model, loss_fn, optimizer):
         optimizer.zero_grad()
 
         if batch % 100 == 0:
-            loss, current = loss.item(), batch * batch_size + len(X)
+            loss, current = loss.item(), batch * len(X)
             print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
 
 
@@ -179,6 +183,7 @@ def test_loop(dataloader, model, loss_fn):
     # also serves to reduce unnecessary gradient computations and memory usage for tensors with requires_grad=True
     with torch.no_grad():
         for X, y in dataloader:
+            X, y = X.to(device), y.to(device)
             pred = model(X)
             test_loss += loss_fn(pred, y).item()
             correct += (pred.argmax(1) == y).type(torch.float).sum().item()