 namespace helpers {
 
     template <typename T>
-    static __global__ void drawBoundignBoxesKernel(T const* images, Nd4jLong* imagesShape, T const* boxes,
-            Nd4jLong* boxesShape, T const* colors, Nd4jLong* colorsShape, T* output, Nd4jLong* outputShape) {
-        for (auto b = 0; b < batchSize; ++b) { // loop by batch
-            for (auto c = 0; c < colorSet->size(); ++c) {
+    static __global__ void drawBoundingBoxesKernel(T const* images, Nd4jLong* imagesShape, T const* boxes,
+            Nd4jLong* boxesShape, T const* colors, Nd4jLong* colorsShape, T* output, Nd4jLong* outputShape,
+            Nd4jLong batchSize, Nd4jLong width, Nd4jLong height, Nd4jLong channels, Nd4jLong colorSetSize) {
+        for (auto b = blockIdx.x; b < batchSize; b += gridDim.x) { // loop by batch
+            for (auto c = 0; c < colorSetSize; c += blockDim.x) {
                 // box with shape
-                auto internalBox = (*boxes)(b, {0})(c, {0});//internalBoxes->at(c);
-                auto color = colorSet->at(c);
-                auto rowStart = nd4j::math::nd4j_max(Nd4jLong (0), Nd4jLong ((height - 1) * internalBox.e<float>(0)));
-                auto rowEnd = nd4j::math::nd4j_min(Nd4jLong (height - 1), Nd4jLong ((height - 1) * internalBox.e<float>(2)));
-                auto colStart = nd4j::math::nd4j_max(Nd4jLong (0), Nd4jLong ((width - 1) * internalBox.e<float>(1)));
-                auto colEnd = nd4j::math::nd4j_min(Nd4jLong(width - 1), Nd4jLong ((width - 1) * internalBox.e<float>(3)));
+                auto pos = 0;
+                auto internalBox = &boxes[pos];//(*boxes)(b, {0})(c, {0});//internalBoxes->at(c);
+                auto color = &colors[pos];//colorSet->at(c);
+                auto rowStart = nd4j::math::nd4j_max(Nd4jLong (0), Nd4jLong ((height - 1) * internalBox[0]));
+                auto rowEnd = nd4j::math::nd4j_min(Nd4jLong (height - 1), Nd4jLong ((height - 1) * internalBox[2]));
+                auto colStart = nd4j::math::nd4j_max(Nd4jLong (0), Nd4jLong ((width - 1) * internalBox[1]));
+                auto colEnd = nd4j::math::nd4j_min(Nd4jLong(width - 1), Nd4jLong ((width - 1) * internalBox[3]));
                 for (auto y = rowStart; y <= rowEnd; y++) {
-                    for (auto e = 0; e < color->lengthOf(); ++e) {
-                        output->p(b, y, colStart, e, color->e(e));
-                        output->p(b, y, colEnd, e, color->e(e));
+                    for (auto e = 0; e < channels; ++e) {
+                        Nd4jLong yMinPos[] = {b, y, colStart, e};
+                        Nd4jLong yMaxPos[] = {b, y, colEnd, e};
+                        auto zIndexYmin = shape::getOffset(outputShape, yMinPos, 4);
+                        auto zIndexYmax = shape::getOffset(outputShape, yMaxPos, 4);
+                        output[zIndexYmin] = color[e];
+                        output[zIndexYmax] = color[e];
                     }
                 }
                 for (auto x = colStart + 1; x < colEnd; x++) {
-                    for (auto e = 0; e < color->lengthOf(); ++e) {
-                        output->p(b, rowStart, x, e, color->e(e));
-                        output->p(b, rowEnd, x, e, color->e(e));
+                    for (auto e = 0; e < channels; ++e) {
+                        Nd4jLong xMinPos[] = {b, rowStart, x, e};
+                        Nd4jLong xMaxPos[] = {b, rowEnd, x, e};
+                        auto zIndexXmin = shape::getOffset(outputShape, xMinPos, 4);
+                        auto zIndexXmax = shape::getOffset(outputShape, xMaxPos, 4);
+                        output[zIndexXmin] = color[e];
+                        output[zIndexXmax] = color[e];
                     }
                 }
             }

    }
    template <typename T>
    void drawBoundingBoxesH(nd4j::LaunchContext* context, NDArray const* images, NDArray const* boxes, NDArray const* colors, NDArray* output) {
        drawBoundingBoxesKernel<<<128, 256, 1024, *stream>>>(imagesBuf, imagesShape, boxesBuf, boxesShape, colorsBuf, colorsShape,
                outputBuf, outputShape, );
    }

    void drawBoundingBoxesFunctor(nd4j::LaunchContext * context, NDArray* images, NDArray* boxes, NDArray* colors, NDArray* output) {

        // colors - colors for each box given
        // set up color for each box as frame

        output->assign(images);
        BUILD_SINGLE_SELECTOR(output->dataType(), drawBoundingBoxesH, (context, images, boxes, colors, output), FLOAT_TYPES);
    }
    BUILD_SINGLE_TEMPLATE(template void drawBoundingBoxesH, (nd4j::LaunchContext* context, NDArray const* images, NDArray const* boxes, NDArray const* colors, NDArray* output), FLOAT_TYPES);
}
}
}
