From eb1f6f52ebbdbf2a39e528a9d63bb91809b4338b Mon Sep 17 00:00:00 2001 From: Jing Guo <3425792+guojing0@users.noreply.github.com> Date: Mon, 1 Sep 2025 15:38:11 +0200 Subject: [PATCH] Enable GPU support; move model and data to GPU if available --- beginner_source/basics/optimization_tutorial.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/beginner_source/basics/optimization_tutorial.py b/beginner_source/basics/optimization_tutorial.py index 82bfaa8f07c..373720945f7 100644 --- a/beginner_source/basics/optimization_tutorial.py +++ b/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()