import torch


def index_initial(n_batch, n_ch, tensor=True):
    """Tensor batch and channel index initialization.

        Args:
            n_batch (Int): Number of batch.
            n_ch (Int): Number of channel.
            tensor (bool): Return tensor or numpy array

        Returns:
            Tensor: Batch index
            Tensor: Channel index

    """
    batch_index = []
    for i in range(n_batch):
        batch_index.append([[i]] * n_ch)
    ch_index = []
    for i in range(n_ch):
        ch_index += [[i]]
    ch_index = [ch_index] * n_batch

    if tensor:
        batch_index = torch.tensor(batch_index)
        ch_index = torch.tensor(ch_index)
        if torch.cuda.is_available():
            batch_index = batch_index.cuda()
            ch_index = ch_index.cuda()
    return batch_index, ch_index