diff --git a/non_square/nvdsinfer_custom_impl_Yolo/yoloPlugins.h b/non_square/nvdsinfer_custom_impl_Yolo/yoloPlugins.h
index 2981af7..177ca10 100644
--- a/non_square/nvdsinfer_custom_impl_Yolo/yoloPlugins.h
+++ b/non_square/nvdsinfer_custom_impl_Yolo/yoloPlugins.h
@@ -56,7 +56,7 @@ class YoloLayer : public nvinfer1::IPluginV2
 {
 public:
     YoloLayer (const void* data, size_t length);
-    YoloLayer (const uint& numBoxes, const uint& numClasses, const uint& gridSize,
+    YoloLayer (const uint& numBoxes, const uint& numClasses, const uint& gridSizeX, const uint& gridSizeY,
                 const uint model_type, const uint new_coords, const float scale_x_y, const float beta_nms,
                 const std::vector<float> anchors, const std::vector<std::vector<int>> mask);
     const char* getPluginType () const override { return YOLOLAYER_PLUGIN_NAME; }
@@ -153,4 +153,4 @@ extern float kBETA_NMS;
 extern std::vector<float> kANCHORS;
 extern std::vector<std::vector<int>> kMASK;
 
-#endif // __YOLO_PLUGINS__
\ No newline at end of file
+#endif // __YOLO_PLUGINS__