     }
     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, );
+        auto batchSize = images->sizeAt(0);
+        auto height = images->sizeAt(1);
+        auto width = images->sizeAt(2);
+        auto channels = images->sizeAt(3);
+        auto stream = context->getCudaStream();
+        auto colorSetSize = colors->sizeAt(0);
+//        auto imageList = images->allTensorsAlongDimension({1, 2, 3}); // split images by batch
+//        auto boxList = boxes->allTensorsAlongDimension({1, 2}); // split boxes by batch
+//        auto colorSet = colors->allTensorsAlongDimension({1});
+        auto imagesBuf = images->getDataBuffer()->specialAsT<T>();
+        auto boxesBuf = boxes->getDataBuffer()->specialAsT<T>();
+        auto colorsBuf = colors->getDataBuffer()->specialAsT<T>();
+        auto outputBuf = output->dataBuffer()->specialAsT<T>();
+        drawBoundingBoxesKernel<<<128, 256, 1024, *stream>>>(imagesBuf, images->getSpecialShapeInfo(),
+                boxesBuf, boxes->getSpecialShapeInfo(), colorsBuf, colors->getSpecialShapeInfo(),
+                outputBuf, output->specialShapeInfo(), batchSize, width, height, channels, colorSetSize);
     }
 
     void drawBoundingBoxesFunctor(nd4j::LaunchContext * context, NDArray* images, NDArray* boxes, NDArray* colors, NDArray* output) {

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) {
                // 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)));
                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 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));
                    }
                }
            }

        // 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);
}
}
}
