diff --git a/src/Algorithm.cpp b/src/Algorithm.cpp
index af12e091..13f8301b 100644
--- a/src/Algorithm.cpp
+++ b/src/Algorithm.cpp
@@ -134,7 +134,7 @@ Algorithm::createParameters()
     KP_LOG_DEBUG("Kompute Algorithm createParameters started");
 
     for (const std::shared_ptr<Memory>& mem : this->mMemObjects) {
-        if (mem->type() == Memory::Type::eImage) {
+        if (mem->getDescriptorType() == vk::DescriptorType::eStorageImage) {
             numImages++;
         } else {
             numTensors++;
diff --git a/src/OpAlgoDispatch.cpp b/src/OpAlgoDispatch.cpp
index c9ba7f55..1f020979 100644
--- a/src/OpAlgoDispatch.cpp
+++ b/src/OpAlgoDispatch.cpp
@@ -25,9 +25,8 @@ OpAlgoDispatch::record(const vk::CommandBuffer& commandBuffer)
 
         // For images the image layout needs to be set to eGeneral before using
         // it for imageLoad/imageStore in a shader.
-        if (mem->type() == Memory::Type::eImage) {
-            std::shared_ptr<Image> image = std::static_pointer_cast<Image>(mem);
-
+        std::shared_ptr<Image> image = std::dynamic_pointer_cast<Image>(mem);
+        if (image) {
             image->recordPrimaryImageBarrier(
               commandBuffer,
               vk::AccessFlagBits::eTransferWrite,