Overload of nonzero is deprecated

[claudiohfg]

Hi, I was executing the example and got the following warning:

[running kmeans]: 0it [00:00, ?it/s]

running k-means on cpu..

/pytorch/torch/csrc/utils/python_arg_parser.cpp:756: UserWarning: This overload of nonzero is deprecated:
	nonzero(Tensor input, *, Tensor out)
Consider using one of the following signatures instead:
	nonzero(Tensor input, *, bool as_tuple)
[running kmeans]: 7it [00:02,  3.24it/s, center_shift=0.000091, iteration=7, tol=0.000100] 

Here is the code I ran to get this message.

import torch
import numpy as np
import matplotlib.pyplot as plt
from kmeans_pytorch import kmeans, kmeans_predict

# set random seed
np.random.seed(123)

# data
data_size, dims, num_clusters = 500000, 100, 3
x = np.random.randn(data_size, dims) / 6
x = torch.from_numpy(x)

# set device
if torch.cuda.is_available():
    device = torch.device('cuda:0')
else:
    device = torch.device('cpu')

# k-means
cluster_ids_x, cluster_centers = kmeans(
    X=x, num_clusters=num_clusters, distance='euclidean', device=device
)

Thanks in advance.

[cbroker1]

yep. same.

[tomerip]

This was solved by PyTorch in this PR: pytorch/vision#2705
A solution that follows this logic without upgrading PyTorch is replacing this line:
selected = torch.nonzero(choice_cluster == index).squeeze().to(device)
with this line:
selected = torch.where((choice_cluster == index) == True)[0].squeeze().to(device)

[sunhao123456789]

怎么解决啊大佬们