def METHOD_NAME(self, idx, n, large_tensor):
    shape = large_tensor.shape
    assert shape[0] % n == 0
    local_shape = [shape[0] // n, shape[1]]
    slice_idx = [
        slice(idx * local_shape[0], (idx + 1) * local_shape[0]),
        slice(local_shape[1]),
    ]
    return large_tensor[slice_idx]