diff --git a/hipamd/src/hip_graph.cpp b/hipamd/src/hip_graph.cpp
index ab72635ecc..e63730aa26 100644
--- a/hipamd/src/hip_graph.cpp
+++ b/hipamd/src/hip_graph.cpp
@@ -171,8 +171,8 @@ hipError_t ihipGraphAddMemcpyNode1D(hip::GraphNode** pGraphNode, hip::Graph* gra
 
 hipError_t ihipGraphAddMemsetNode(hip::GraphNode** pGraphNode, hip::Graph* graph,
                                   hip::GraphNode* const* pDependencies, size_t numDependencies,
-                                  const hipMemsetParams* pMemsetParams,
-                                  bool capture = true, size_t depth = 1) {
+                                  const hipMemsetParams* pMemsetParams, bool capture = true,
+                                  size_t depth = 1, size_t arrWidth = 1, size_t arrHeight = 1) {
   if (pGraphNode == nullptr || graph == nullptr || pMemsetParams == nullptr ||
       (numDependencies > 0 && pDependencies == nullptr) || pMemsetParams->height == 0) {
     return hipErrorInvalidValue;
@@ -208,7 +208,7 @@ hipError_t ihipGraphAddMemsetNode(hip::GraphNode** pGraphNode, hip::Graph* graph
   if (status != hipSuccess) {
     return status;
   }
-  *pGraphNode = new hip::GraphMemsetNode(pMemsetParams, depth);
+  *pGraphNode = new hip::GraphMemsetNode(pMemsetParams, depth, arrWidth, arrHeight);
   status = ihipGraphAddNode(*pGraphNode, graph, pDependencies, numDependencies, capture);
   return status;
 }
@@ -747,10 +747,16 @@ hipError_t capturehipMemset3DAsync(hipStream_t& stream, hipPitchedPtr& pitchedDe
                                    hipExtent& extent) {
   ClPrint(amd::LOG_INFO, amd::LOG_API, "[hipGraph] Current capture node Memset3D on stream : %p",
           stream);
-  hipMemsetParams memsetParams = {0};
   if (!hip::isValid(stream)) {
     return hipErrorContextIsDestroyed;
   }
+
+  // In this case no work should be done
+  if (extent.width == 0 || extent.height == 0 || extent.depth == 0) {
+    return hipSuccess;
+  }
+
+  hipMemsetParams memsetParams = {0};
   memsetParams.dst = pitchedDevPtr.ptr;
   memsetParams.value = value;
   memsetParams.width = extent.width;
@@ -761,7 +767,8 @@ hipError_t capturehipMemset3DAsync(hipStream_t& stream, hipPitchedPtr& pitchedDe
   hip::GraphNode* pGraphNode;
   hipError_t status =
       ihipGraphAddMemsetNode(&pGraphNode, s->GetCaptureGraph(), s->GetLastCapturedNodes().data(),
-                             s->GetLastCapturedNodes().size(), &memsetParams, true, extent.depth);
+                             s->GetLastCapturedNodes().size(), &memsetParams, true, extent.depth,
+                             pitchedDevPtr.xsize, pitchedDevPtr.ysize);
   if (status != hipSuccess) {
     return status;
   }
diff --git a/hipamd/src/hip_graph_internal.hpp b/hipamd/src/hip_graph_internal.hpp
index a75d431a3e..db5256f75b 100644
--- a/hipamd/src/hip_graph_internal.hpp
+++ b/hipamd/src/hip_graph_internal.hpp
@@ -1979,11 +1979,16 @@ class GraphMemcpyNodeToSymbol : public GraphMemcpyNode1D {
 class GraphMemsetNode : public GraphNode {
   hipMemsetParams memsetParams_;
   size_t depth_ = 1;
+  size_t arrWidth_ = 1;
+  size_t arrHeight_ = 1;
  public:
-  GraphMemsetNode(const hipMemsetParams* pMemsetParams, size_t depth = 1)
+  GraphMemsetNode(const hipMemsetParams* pMemsetParams, size_t depth = 1, size_t arrWidth = 1,
+                  size_t arrHeight = 1)
       : GraphNode(hipGraphNodeTypeMemset, "solid", "invtrapezium", "MEMSET") {
     memsetParams_ = *pMemsetParams;
     depth_ = depth;
+    arrWidth_ = arrWidth;
+    arrHeight_ = arrHeight;
     size_t sizeBytes = 0;
     if (memsetParams_.height == 1) {
       sizeBytes = memsetParams_.width * memsetParams_.elementSize;
@@ -1997,6 +2002,8 @@ class GraphMemsetNode : public GraphNode {
   GraphMemsetNode(const GraphMemsetNode& memsetNode) : GraphNode(memsetNode) {
     memsetParams_ = memsetNode.memsetParams_;
     depth_ = memsetNode.depth_;
+    arrWidth_ = memsetNode.arrWidth_;
+    arrHeight_ = memsetNode.arrHeight_;
   }
 
   GraphNode* clone() const override {
@@ -2040,15 +2047,15 @@ class GraphMemsetNode : public GraphNode {
     if (status != hipSuccess) {
       return status;
     }
-    if (memsetParams_.height == 1) {
+    if (memsetParams_.height == 1 && depth_ == 1) {
       size_t sizeBytes = memsetParams_.width * memsetParams_.elementSize;
       hipError_t status = ihipMemsetCommand(commands_, memsetParams_.dst, memsetParams_.value,
                                             memsetParams_.elementSize, sizeBytes, stream);
     } else {
       hipError_t status = ihipMemset3DCommand(
           commands_,
-          {memsetParams_.dst, memsetParams_.pitch, memsetParams_.width * memsetParams_.elementSize,
-           memsetParams_.height},
+          {memsetParams_.dst, memsetParams_.pitch, arrWidth_ * memsetParams_.elementSize,
+           arrHeight_},
           memsetParams_.value,
           {memsetParams_.width * memsetParams_.elementSize, memsetParams_.height, depth_}, stream,
           memsetParams_.elementSize);