import torch

def stack_jagged(tensors, fill_value = 0):
    # does not use F.pad to not allocate, could as well use F.pad's arg names instead 
    shape = [len(tensors)] + [max(t.shape[dim] for t in tensors) for dim in range(len(tensors[0].shape))]
    res = torch.full(shape, fill_value, dtype = tensors[0].dtype, device = tensors[0].device)
    for r, t in zip(res, tensors):
        r[tuple(map(slice, t.shape))] = t
    return res

if __name__ == '__main__':
    print(stack_jagged([torch.zeros(3, 4, 5), torch.zeros(3, 5, 6)]).shape)