def METHOD_NAME(func: Callable):
    r"""A decorator to test both on CPU and CUDA (if available)."""
    import pytest
    devices = [torch.device('cpu')]
    if torch.cuda.is_available():
        devices.append(torch.device('cuda:0'))
    return pytest.mark.parametrize('device', devices)(func)