INNER CODE UNIT · Python

sources

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

    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)

    segment_arange_cuda_module = module

    return segment_arange_cuda_module


def segment_arange(
    ranges: torch.Tensor,
    start: int,
    seq_len: int,

View source record →