File tree Expand file tree Collapse file tree 1 file changed +4
-4
lines changed Expand file tree Collapse file tree 1 file changed +4
-4
lines changed Original file line number Diff line number Diff line change @@ -5,10 +5,10 @@ at::Tensor nms_cpu_kernel(
55 const at::Tensor& dets,
66 const at::Tensor& scores,
77 const float iou_threshold) {
8- AT_ASSERTM (!dets.type ().is_cuda (), " dets must be a CPU tensor" );
9- AT_ASSERTM (!scores.type ().is_cuda (), " scores must be a CPU tensor" );
8+ AT_ASSERTM (!dets.options (). device ().is_cuda (), " dets must be a CPU tensor" );
9+ AT_ASSERTM (!scores.options (). device ().is_cuda (), " scores must be a CPU tensor" );
1010 AT_ASSERTM (
11- dets.type () == scores.type (), " dets should have the same type as scores" );
11+ dets.scalar_type () == scores.scalar_type (), " dets should have the same type as scores" );
1212
1313 if (dets.numel () == 0 )
1414 return at::empty ({0 }, dets.options ().dtype (at::kLong ));
@@ -74,7 +74,7 @@ at::Tensor nms_cpu(
7474 const float iou_threshold) {
7575 auto result = at::empty ({0 }, dets.options ());
7676
77- AT_DISPATCH_FLOATING_TYPES (dets.type (), " nms" , [&] {
77+ AT_DISPATCH_FLOATING_TYPES (dets.scalar_type (), " nms" , [&] {
7878 result = nms_cpu_kernel<scalar_t >(dets, scores, iou_threshold);
7979 });
8080 return result;
You can’t perform that action at this time.
0 commit comments