Workaround for INDArray.get rank1/2 issue (#7092)

deeplearning4j · Feb 5, 2019 · 5d63924 · 5d63924
1 parent 5e211ba
commit 5d63924
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/...i-parent/nd4j-api/src/main/java/org/nd4j/autodiff/samediff/internal/InferenceSession.java b/...i-parent/nd4j-api/src/main/java/org/nd4j/autodiff/samediff/internal/InferenceSession.java
@@ -335,6 +335,10 @@ public INDArray[] getOutputs(DifferentialFunction op, FrameIter outputFrameIter,
                     idx[0] = NDArrayIndex.point(i);
                     INDArray get = valuesArr.get(idx).dup();
                     int outIdx = idxs[i];
+                    if(idxs.length == 2 && get.rank() == 2){
+                        //Workaround for: https://github.com/deeplearning4j/deeplearning4j/issues/7092
+                        get = get.reshape(get.length());
+                    }
                     l.set(outIdx, get);
                 }