import torch


def unpack_bits(t):
    """t is a tensor of packed bits. Unpack in the first dim"""
    n = t.shape[0]
    unpacked = torch.zeros(8*n, *t.shape[1:]).to(t.device)
    for i in range(8):
        unpacked[i*n:(i+1)*n] = (t >> i) & 1
    return unpacked