RuntimeError: Could not run 'aten::nonzero' with arguments from the 'SparseCUDA' backend. 'aten::nonzero' is only available for these backends: [CPU, CUDA, Autograd, Profiler, Tracer].

[derronqi]

pytorch1.6:
cuda10.2
titan rtx * 4

output = self.am_branches[i](x.cuda(i), labels[i])

File "/home/derron/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/derron/arcface-pytorch/head/metrics_parallel.py", line 102, in forward
output[index] = phi[index]
RuntimeError: Could not run 'aten::nonzero' with arguments from the 'SparseCUDA' backend. 'aten::nonzero' is only available for these backends: [CPU, CUDA, Autograd, Profiler, Tracer].