sun_diffusion.devices module¶
Utilities for handling PyTorch devices and dtypes.
- sun_diffusion.devices.set_device(device=None, cuda_id=0)¶
Set global device and default dtype for torch.
- Parameters:
device (str) – ‘cpu’ or ‘cuda’. If None, defaults to cuda if available
cuda_id (int) – Which CUDA device to use
- Return type:
None
- sun_diffusion.devices.get_device()¶
Return the current torch.device.
- Return type:
device
- sun_diffusion.devices.get_dtype(is_complex=False)¶
Return the current default dtype (complex or real).
- Return type:
dtype
- sun_diffusion.devices.set_dtype(dtype)¶
Sets the default dtype to dtype.
- Return type:
None
- sun_diffusion.devices.device_name()¶
Get human-readable name for current device.
- Return type:
str
- sun_diffusion.devices.summary()¶
Returns a summary string of the current device setup.
- Return type:
str