#6358 Fix issue with LayerWorkspaceMgr when used directly in conjunct…

…ion with CuDNN
deeplearning4j · Sep 5, 2018 · 0c21d7c · 0c21d7c
1 parent a077937
commit 0c21d7c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/...4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/workspace/LayerWorkspaceMgr.java b/...4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/workspace/LayerWorkspaceMgr.java
@@ -96,7 +96,7 @@ public INDArray validateArrayLocation(@NonNull ArrayType arrayType, @NonNull IND
      * @return Pointer for that key, or null if none exists
      */
     public <T extends Pointer> T getHelperWorkspace(String key){
-        return (T)helperWorkspacePointers.get(key);
+        return helperWorkspacePointers == null ? null : (T)helperWorkspacePointers.get(key);
     }
 
     /**
@@ -108,6 +108,9 @@ public <T extends Pointer> T getHelperWorkspace(String key){
      * @param value Pointer
      */
     public void setHelperWorkspace(@NonNull String key, Pointer value){
+        if(helperWorkspacePointers == null){
+            helperWorkspacePointers = new HashMap<>();
+        }
         helperWorkspacePointers.put(key, value);
     }