model.to(device)
时间: 2024-05-12 08:20:51 浏览: 53
This line of code is used to move a PyTorch model to a specified device, such as a GPU or CPU. By default, PyTorch models are created on the CPU. However, if the user has a GPU available, they can move the model to the GPU to take advantage of its parallel processing capabilities and speed up the model's training and inference.
The "device" parameter is a PyTorch device object that can be created using the `torch.device()` function. For example, if the user wants to move the model to the first available GPU, they can use the following code:
```
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
```
This code will check if a GPU is available and set the device to "cuda:0" (the first available GPU) if it is, or "cpu" if not. The model will then be moved to the specified device.
阅读全文