19 lines
525 B
Python
19 lines
525 B
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tools.lab_extension import build_extension
|
|
|
|
|
|
@pytest.mark.cuda_required
|
|
@pytest.mark.skeleton
|
|
def test_extension_can_build_or_skip():
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("CUDA is not available")
|
|
ext = build_extension(verbose=False)
|
|
if ext is None:
|
|
pytest.skip("extension build/load is unavailable in this environment")
|
|
assert hasattr(torch.ops, "kernel_lab")
|
|
assert hasattr(torch.ops.kernel_lab, "vector_add")
|