Hello,
I am doing the Udemy course on Pytorch and trying to run the MNIST example on the GPU. Even though I activated CUDA, training is slow and GPU usage is at 0%.
Here is my code :
device = 'cuda'
model.to(device);

# number of epochs to train the model
n_epochs = 30  # suggest training between 20-50 epochs


model.train() # prep model for training

for epoch in range(n_epochs):
    # monitor training loss
    train_loss = 0.0
    
    ###################
    # train the model #
    ###################
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        start = time.time()
        # clear the gradients of all optimized variables
        optimizer.zero_grad()
        # forward pass: compute predicted outputs by passing inputs to the model
        output = model(data)
        # calculate the loss
        loss = criterion(output, target)
        # backward pass: compute gradient of the loss with respect to model parameters
        loss.backward()
        # perform a single optimization step (parameter update)
        optimizer.step()

        # update running training loss
        train_loss += loss.item()*data.size(0)
    print(f"Device = {device}; Time per batch: {(time.time() - start)/3:.3f} seconds")
    print(' ')    
    # print training statistics 
    # calculate average loss over an epoch
    train_loss = train_loss/len(train_loader.sampler)

    print('Epoch: {} \tTraining Loss: {:.6f}'.format(
        epoch+1, 
        train_loss
        ))
    print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB')
    print('Cached:   ', round(torch.cuda.memory_cached(0)/1024**3,1), 'GB')
    print(' ')

Should I add other lines of code? Thank you