9a2903dac097a7f3633f6c828533ed385c7159c9,qucumber/utils/data.py,,extract_refbasis_samples,#,64

Before Change

    
    tmp = []
    num_visible = train_samples.shape[-1]
    for i in range(train_samples.shape[0]):
        flag = 0
        for j in range(num_visible):
            if train_bases[i][j] != "Z":
                flag = 1
                break
        if flag == 0:
            tmp.append(train_samples[i])
    z_samples = torch.zeros(len(tmp), num_visible, dtype=torch.double)
    for i in range(len(tmp)):
        for j in range(num_visible):
            z_samples[i][j] = tmp[i][j]

After Change

    :returns: The samples in the data that are only in the reference basis.
    :rtype: torch.Tensor
    
    idx = (
        torch.tensor(train_bases == "Z", dtype=torch.uint8)
        .all(dim=1)
        .to(train_samples.device)
    )
    z_samples = train_samples[idx]
    return z_samples

Instances