pytorch 批量篩選

 

這個是按行篩選

    import torch

    input_tensor = torch.tensor([[1, 2, 3],
                                 [4, 5,6]])
    # print(input_tensor > torch.tensor([2,3,4]))
    # tensor([[False, False, False],
    #         [True, True, True]])

    print(input_tensor)
    ddd=input_tensor > torch.tensor([[1,2,2],[2,2,2]])
    print(input_tensor[ddd])