Fix nonzero bug

Filip Binkiewicz · apaszke · commit 9510431716e8 · 2017-04-28T16:10:30.000-07:00
diff --git a/torch/lib/THD/master_worker/master/generic/THDTensorMath.cpp b/torch/lib/THD/master_worker/master/generic/THDTensorMath.cpp
@@ -70,7 +70,7 @@ void THDTensor_(nonzero)(THDLongTensor *subscript, THDTensor *tensor) {
     ),
     THDState::s_current_worker
   );
-  ptrdiff_t numel = receiveValueFromWorker<ptrdiff_t>(tensor->storage->node_id);
+  long long numel = receiveValueFromWorker<long long>(tensor->storage->node_id);
   THDLongTensor__resize2d(subscript, numel, tensor->nDimension);
 }
 
diff --git a/torch/lib/THD/master_worker/worker/dispatch/TensorMath.cpp b/torch/lib/THD/master_worker/worker/dispatch/TensorMath.cpp
@@ -53,7 +53,8 @@ static void tensorNonzero(rpc::RPCMessage& raw_message) {
   thpp::Tensor *tensor = unpackRetrieveTensor(raw_message);
   finalize(raw_message);
   tensor->nonzero(*subscript);
-  sendValueToMaster((double)subscript->numel());
+  long long numel = subscript->sizes().size() > 0 ? subscript->sizes()[0] : 0;
+  sendValueToMaster(numel);
 }
 
 static void tensorIndexSelect(rpc::RPCMessage& raw_message) {

Original file line number	Diff line number	Diff line change
`@@ -70,7 +70,7 @@ void THDTensor_(nonzero)(THDLongTensor subscript, THDTensor tensor) {`
`70`	`70`	`),`
`71`	`71`	`THDState::s_current_worker`
`72`	`72`	`);`
`73`		`- ptrdiff_t numel = receiveValueFromWorker<ptrdiff_t>(tensor->storage->node_id);`
	`73`	`+ long long numel = receiveValueFromWorker<long long>(tensor->storage->node_id);`
`74`	`74`	`THDLongTensor__resize2d(subscript, numel, tensor->nDimension);`
`75`	`75`	`}`
`76`	`76`
Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,8 @@ static void tensorNonzero(rpc::RPCMessage& raw_message) {`
`53`	`53`	`thpp::Tensor *tensor = unpackRetrieveTensor(raw_message);`
`54`	`54`	`finalize(raw_message);`
`55`	`55`	`tensor->nonzero(*subscript);`
`56`		`- sendValueToMaster((double)subscript->numel());`
	`56`	`+ long long numel = subscript->sizes().size() > 0 ? subscript->sizes()[0] : 0;`
	`57`	`+ sendValueToMaster(numel);`
`57`	`58`	`}`
`58`	`59`
`59`	`60`	`static void tensorIndexSelect(rpc::RPCMessage& raw_message) {`