Updated reset

jungleTail · Mar 9, 2016 · 0818cb9 · 0818cb9
1 parent ab69351
commit 0818cb9
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/...4j-core/src/main/java/org/deeplearning4j/datasets/canova/RecordReaderDataSetIterator.java b/...4j-core/src/main/java/org/deeplearning4j/datasets/canova/RecordReaderDataSetIterator.java
@@ -242,8 +242,11 @@ public int totalOutcomes() {
 
     @Override
     public void reset() {
-        if (recordReader instanceof RecordReader)
+        if (recordReader instanceof RecordReader){
+            batchNum = 0;
+            notOvershot = true;
             recordReader.reset();
+        }
         else if (recordReader instanceof SequenceRecordReader)
             throw new UnsupportedOperationException("Reset not supported for SequenceRecordReader type.");
     }