[Device] managing device

[hanson-hschang]

some idea of managing device, also see torch.set_default_device

import torch

from ss.utility.logging import Logging

logger = Logging.get_logger(__name__)


class Device:

    @staticmethod
    def get() -> torch.device:
    
        """Get the best available device."""
        # Check if CUDA is available
        if torch.cuda.is_available():
            logger.debug("Using CUDA device.")
            return torch.device("cuda")

        # Check if MPS (Apple M series chip) is available
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            logger.debug("Using MPS device.")
            return torch.device("mps")

        # If no CUDA or MPS device is available, use CPU
        logger.debug("Using CPU device.")
        return torch.device("cpu")

in addition, the mps device might be slower than cpu? check this issue.

[hanson-hschang]

Distributed Data Parallel in PyTorch

[hanson-hschang]

• Useful tutorial on how to train with multiple GPU: https://www.youtube.com/watch?v=TibQO_xv1zc