YouTipPytorch Torch Set_Default_Device
## PyTorch: torch.set_default_device
In PyTorch, managing where your tensors are allocated (CPU vs. GPU) is crucial for performance. Traditionally, you had to explicitly pass the `device` argument (e.g., `device='cuda'`) to every tensor creation factory function or call `.to('cuda')` on existing tensors.
Introduced in PyTorch 2.0, `torch.set_default_device` simplifies this workflow by allowing you to set a global default device. Any newly created tensors will automatically be allocated on this specified device unless an explicit device is provided.
---
### Function Definition
```python
torch.set_default_device(device)
```
#### Parameters:
* **`device`** *(torch.device or str)*: The destination device (e.g., `'cpu'`, `'cuda'`, `'cuda:0'`, or `'mps'`) where newly created tensors should be allocated by default. Passing `None` resets the default device to the standard CPU device.
---
### Code Example
The following example demonstrates how to set the default device to GPU (CUDA) if available, create a tensor without specifying a device, and then reset the default device back to CPU.
```python
import torch
# Check if CUDA (GPU) is available
if torch.cuda.is_available():
# Set the default device to CUDA
torch.set_default_device('cuda')
# This tensor is automatically created on the GPU (CUDA)
x = torch.zeros(3, 4)
print("Default device tensor location:", x.device)
# Reset the default device back to CPU
torch.set_default_device('cpu')
# This tensor is created on the CPU
y = torch.zeros(3, 4)
print("Reset device tensor location:", y.device)
else:
print("CUDA is not available. Defaulting to CPU.")
```
---
### Key Considerations & Best Practices
1. **Factory Functions Only**: `torch.set_default_device` affects PyTorch tensor factory functions (such as `torch.arange`, `torch.empty`, `torch.ones`, `torch.zeros`, `torch.rand`, etc.) and neural network module initializations.
2. **Explicit Overrides**: You can still override the default device at any time by passing an explicit `device` argument to a factory function (e.g., `torch.zeros(3, 4, device='cpu')` will still be created on the CPU even if the default is set to `'cuda'`).
3. **Context Manager Alternative**: If you only want to temporarily change the default device for a specific block of code, it is highly recommended to use the context manager `with torch.device(device):` instead of globally changing it with `torch.set_default_device`. This prevents side effects in other parts of your codebase.
#### Example using the Context Manager:
```python
# Temporarily allocate tensors on CUDA
with torch.device('cuda'):
# x is created on CUDA
x = torch.ones(2, 2)
# y is created on the global default device (CPU)
y = torch.ones(2, 2)
```