INNER CODE UNIT · Python

load_cuda_module

microsoft/muzic · museformer/museformer/kernels/segment_arange/main.py:12

def load_cuda_module():
    global segment_arange_cuda_module

    if segment_arange_cuda_module is not None:
        return segment_arange_cuda_module

    path = os.path.dirname(os.path.abspath(__file__))
    cpps = glob.glob(os.path.join(path, "cuda_src/*.cpp"))
    cudas = glob.glob(os.path.join(path, "cuda_src/*.cu"))
    sources = list(cpps) + list(cudas)

    from torch.utils.cpp_extension import load
    module = load(name='segment_arange_cuda',
                  sources=sources,
                  extra_cflags=['-O2'],
                  with_cuda=True,
                  verbose=False)

View source record →