Update custom_layernorm.cu

NVIDIA-AI-IOT · Dec 3, 2024 · 8a1b296 · 8a1b296
1 parent b248734
commit 8a1b296
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/CUDA-BEVFusion/src/plugins/custom_layernorm.cu b/CUDA-BEVFusion/src/plugins/custom_layernorm.cu
@@ -121,7 +121,12 @@ public:
 
     virtual int32_t enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
         void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept{
-
+
+        if(inputDesc[0].dims.nbDims != 3){
+            printf("Unsupported tensor dims: %d (expected 3)\n", inputDesc[0].dims.nbDims);
+            return 1;
+        }
+
         // B, N, C
         int N = inputDesc[0].dims.d[0] * inputDesc[0].dims.d[1];
         int C = inputDesc[0].dims.d[2];
@@ -141,6 +146,12 @@ public:
             // not implemented
             return 1;
         }
+
+        auto code = cudaPeekAtLastError();
+        if(code != cudaSuccess){
+            printf("Failed to run kernel(layernorm_kernel) with dtype %d\n", (int)inputDesc[0].type);
+            return 1;
+        }
         return 0;
     }