Course project using GPU acceleration

I am walking through the Jetson Nano AI course, and was using the nvdli-nano to run the CNN on Jetson Nano. I went through the code lines in the jupyter notebook, and don’t find a line that specify the training to be performed in GPU. I wonder if that is inferred somewhere, or set by default? If I have both a CPU and GPU, how should I allocate the computational power of each to perform the task?

Hi,

Please noted that Jetson is designed mainly for inference.
For training on Jetson, you can check if this page can meet your requirement:

To check if a framework is running on GPU, you can use the API like this:

import torch
print(torch.cuda.is_available())

Thanks.

In the nvdli-nano notebooks, if you look at where the model is initially created, there are these lines of code:

device = torch.device('cuda')

# model is created...

model = model.to(device)

This tells PyTorch to run the model on the CUDA device, and hence both training and inference will be done using the GPU.