set default workspaceIndexMode to image index for 3D pooling (#366)

ROCm · Aug 8, 2020 · c16087a · c16087a
1 parent af79ca7
commit c16087a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/pooling.cpp b/src/pooling.cpp
@@ -57,7 +57,8 @@ PoolingDescriptor::PoolingDescriptor(miopenPoolingMode_t m,
       mode(m),
       pmode(pm),
       indexType(miopenIndexUint8),
-      workspaceIndexMode(miopenPoolingWorkspaceIndexMask)
+      workspaceIndexMode(size == 3 ? miopenPoolingWorkspaceIndexImage
+                                   : miopenPoolingWorkspaceIndexMask)
 {
 }
 
@@ -74,6 +75,8 @@ PoolingDescriptor::PoolingDescriptor(miopenPoolingMode_t m,
       indexType(miopenIndexUint8),
       workspaceIndexMode(miopenPoolingWorkspaceIndexMask)
 {
+    if(plens.size() == 3)
+        workspaceIndexMode = miopenPoolingWorkspaceIndexImage;
 }
 
 void PoolingDescriptor::SetIndexType(miopenIndexType_t index_type) { indexType = index_type; }