improve the code

Wrench-Git · Wrench-Git · commit ad2e7e0b5a41 · 2024-07-11T13:39:44.000+08:00
diff --git a/dipu/torch_dipu/csrc_dipu/aten/ops/OpUtils.hpp b/dipu/torch_dipu/csrc_dipu/aten/ops/OpUtils.hpp
@@ -254,8 +254,8 @@ inline bool is_scalar_on_cpu(const at::Tensor& t) {
   return t.defined() && t.is_cpu() && t.numel() == 1;
 }
 
-inline bool check_scalar_on_cpu(const c10::optional<at::Tensor> t) {
-  return t.has_value() && (*t).unsafeGetTensorImpl()->is_wrapped_number();
+inline bool check_scalar_on_cpu(const c10::optional<at::Tensor>& t) {
+  return t.has_value() && ((*t).unsafeGetTensorImpl()->is_wrapped_number() || ((*t).is_cpu() && (*t).numel() == 1 ));
 }
 
 }  // namespace native

Original file line number	Diff line number	Diff line change
`@@ -254,8 +254,8 @@ inline bool is_scalar_on_cpu(const at::Tensor& t) {`
`254`	`254`	`return t.defined() && t.is_cpu() && t.numel() == 1;`
`255`	`255`	`}`
`256`	`256`
`257`		`-inline bool check_scalar_on_cpu(const c10::optional<at::Tensor> t) {`
`258`		`- return t.has_value() && (*t).unsafeGetTensorImpl()->is_wrapped_number();`
	`257`	`+inline bool check_scalar_on_cpu(const c10::optional<at::Tensor>& t) {`
	`258`	`+ return t.has_value() && ((t).unsafeGetTensorImpl()->is_wrapped_number() \|\| ((t).is_cpu() && (*t).numel() == 1 ));`
`259`	`259`	`}`
`260`	`260`
`261`	`261`	`} // namespace native`