[Feature](mluOpAdamW): add param check for adam_w (#984)

Cambricon · Apr 2, 2024 · 94fbbc0 · 94fbbc0
1 parent 0d745d4
commit 94fbbc0
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/kernels/adam_w/adam_w.cpp b/kernels/adam_w/adam_w.cpp
@@ -97,6 +97,17 @@ mluOpAdamW(mluOpHandle_t handle, const mluOpAdamWDescriptor_t adamw_desc,
   PARAM_CHECK("[mluOpAdamW]", momentum_desc != nullptr);
   PARAM_CHECK("[mluOpAdamW]", velocity_desc != nullptr);
   PARAM_CHECK("[mluOpAdamW]", grad_desc != nullptr);
+  PARAM_CHECK("[mluOpAdamW]", param_desc->dtype == MLUOP_DTYPE_FLOAT)
+  PARAM_CHECK("[mluOpAdamW]", paramh_desc->dtype == MLUOP_DTYPE_BFLOAT16)
+  PARAM_CHECK("[mluOpAdamW]", momentum_desc->dtype == MLUOP_DTYPE_FLOAT)
+  PARAM_CHECK("[mluOpAdamW]", velocity_desc->dtype == MLUOP_DTYPE_FLOAT)
+  PARAM_CHECK("[mluOpAdamW]", grad_desc->dtype == MLUOP_DTYPE_BFLOAT16)
+
+  PARAM_CHECK_LE("[mluOpAdamW]", beta1, 1.0)
+  PARAM_CHECK_GE("[mluOpAdamW]", beta1, 0.0)
+  PARAM_CHECK_LE("[mluOpAdamW]", beta2, 1.0)
+  PARAM_CHECK_GE("[mluOpAdamW]", beta2, 0.0)
+  PARAM_CHECK("[mluOpAdamW]", epsilon > 0)
 
   size_t param_dims = 0;
   size_t paramh_dims = 0;