pytorch - 为什么 torch::nonzero() 非常慢？

Question

我执行了下面的代码。但是 nonzero() 函数非常慢（需要 0.1 秒）。我的代码有什么问题？

torch::Tensor a = torch::randn({30000, 80});
torch::Tensor b = a.select(1, 4) > 0.2;
const torch::Tensor c = torch::nonzero(b);
.
.
.

Answer 1

你的代码没有错。torch::nonzero在这种情况下，必须检查您给它的张量中的每个条目30000*80=2.4e+6。这只是需要一段时间。