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(
5
5
const at::Tensor& dets,
6
6
const at::Tensor& scores,
7
7
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" );
10
10
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" );
12
12
13
13
if (dets.numel () == 0 )
14
14
return at::empty ({0 }, dets.options ().dtype (at::kLong ));
@@ -74,7 +74,7 @@ at::Tensor nms_cpu(
74
74
const float iou_threshold) {
75
75
auto result = at::empty ({0 }, dets.options ());
76
76
77
- AT_DISPATCH_FLOATING_TYPES (dets.type (), " nms" , [&] {
77
+ AT_DISPATCH_FLOATING_TYPES (dets.scalar_type (), " nms" , [&] {
78
78
result = nms_cpu_kernel<scalar_t >(dets, scores, iou_threshold);
79
79
});
80
80
return result;
You can’t perform that action at this time.
0 commit comments