2c01aa22c697b4326db05b7729f85613ebb8a8a0,torch_geometric/utils/one_hot.py,,one_hot,#Any#Any#Any#,4

Before Change

    if num_classes is None:
        num_classes = src.max(dim=0)[0] + 1
    elif isinstance(num_classes, int) or isinstance(num_classes, float):
        num_classes = torch.tensor(int(num_classes))

    if src.size(1) > 1:
        zero = torch.tensor([0], device=src.device)
        src = src + torch.cat([zero, torch.cumsum(num_classes, 0)[:-1]])

After Change

    if num_classes is None:
        num_classes = src.max(dim=0)[0] + 1
    else:
        num_classes = torch.tensor(
            repeat(num_classes, length=src.size(1)),
            dtype=torch.long,
            device=src.device)

    if src.size(1) > 1:
        zero = torch.tensor([0], device=src.device)
        src = src + torch.cat([zero, torch.cumsum(num_classes, 0)[:-1]])

Instances