INNER CODE UNIT · Python
segment_arange_pytorch
microsoft/muzic · museformer/museformer/kernels/segment_arange/main.py:76
def segment_arange_pytorch(ranges, start, seq_len, pad_value, dtype=torch.long):
out = torch.full((seq_len,), pad_value, dtype=dtype, device=ranges.device)
for idx, (begin, end) in enumerate(ranges):
out[begin: end] = torch.arange(start, start + (end - begin), dtype=dtype, device=ranges.device)
return out