PyTorch torch.cuda Functions
Pytorch torch Reference Manual
torch.cudaIt is a module in PyTorch for CUDA management. It provides functions such as CUDA device management, memory management, and synchronization.
Function Definition
torch.cuda
The torch.cuda module contains the following commonly used functions:
torch.cuda.is_available()- Check if CUDA is availabletorch.cuda.device_count()- Get the number of CUDA devicestorch.cuda.current_device()- Get the current devicetorch.cuda.synchronize()- Synchronize CUDA operationstorch.cuda.empty_cache()- Clear cache
Usage Example
Example
import torch
# Check if CUDA is available
print(f"CUDA available: {torch.cuda.is_available()}")
# Get the number of CUDA devices
if torch.cuda.is_available():
print(f"CUDA device count: {torch.cuda.device_count()}")
print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
# Check if CUDA is available
print(f"CUDA available: {torch.cuda.is_available()}")
# Get the number of CUDA devices
if torch.cuda.is_available():
print(f"CUDA device count: {torch.cuda.device_count()}")
print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
Other Extensions