import torch


def get_batches(s_t, b=None, device=None):
    """
    Get batches from a 3d sparse tensor.
    :param s_t: sparse tensor
    :param b: if None, return all batches in a list; else, return a specific batch
    :param device: device. If None, it will be the same as input
    :return: sparse tensor or list of sparse tensors
    """
    if device is None:
        device = s_t.device

    coo = s_t._indices()
    data = s_t._values()
    if b is not None:
        idx = (coo[0, :] == b).nonzero()
        _coo = coo[1:3, idx].view(2, -1)
        _data = data[idx].view(-1)
        outp = torch.sparse_coo_tensor(_coo, _data, s_t.shape[1:3], dtype=_data.dtype, device=device)
    else:
        batch_num = s_t.shape[0]
        outp = []
        for b in range(batch_num):
            idx = (coo[0, :] == b).nonzero()
            _coo = coo[1:3, idx].view(2, -1)
            _data = data[idx].view(-1)
            outp.append(torch.sparse_coo_tensor(_coo, _data, s_t.shape[1:3], dtype=_data.dtype, device=device))
    return outp