Skip to content

framework

zeus.util.framework

Utilities for framework-specific code.

torch_is_available cached 

torch_is_available()

Check if PyTorch is available.

Source code in zeus/util/framework.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
@lru_cache(maxsize=1)
def torch_is_available():
    """Check if PyTorch is available."""
    try:
        import torch

        assert (
            torch.cuda.is_available()
        ), "PyTorch is available but does not have CUDA support."
        MODULE_CACHE["torch"] = torch
        logger.info("PyTorch with CUDA support is available.")
        return True
    except ImportError:
        logger.info("PyTorch is not available.")
        return False

cuda_sync 

cuda_sync(device=None)

Synchronize CPU and CUDA.

cupy.cuda.Device.synchronize may be a good choice to make

CUDA device synchronization more general. Haven't tested it yet.

Parameters:

Name Type Description Default
device int | None

The device to synchronize.

 None
Source code in zeus/util/framework.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def cuda_sync(device: int | None = None) -> None:
    """Synchronize CPU and CUDA.

    Note: `cupy.cuda.Device.synchronize` may be a good choice to make
          CUDA device synchronization more general. Haven't tested it yet.

    Args:
        device: The device to synchronize.
    """
    if torch_is_available():
        torch = MODULE_CACHE["torch"]
        torch.cuda.synchronize(device)
        return

    raise RuntimeError("No frameworks are available.")