Fixed the infinite set flag.

2 files changed, 4 insertions, 5 deletions

diff --git a/DataSetLabelMe.lua b/DataSetLabelMe.lua
index 8d81448..05ed8eb 100644
--- a/DataSetLabelMe.lua
+++ b/DataSetLabelMe.lua
@@ -197,7 +197,6 @@ function DataSetLabelMe:__index__(key)
             tag_idx = math.floor((key-1)/self.nbClasses) + 1
          end
          tag_idx = ((tag_idx-1) % (self.tags[ctr_target].size/3))*3 + 1
-         print('key:', key, 'tag:', tag_idx, 'label', ctr_target)
       end
 
       -- generate patch
diff --git a/StochasticTrainer.lua b/StochasticTrainer.lua
index 8c71f3e..5a369c4 100644
--- a/StochasticTrainer.lua
+++ b/StochasticTrainer.lua
@@ -166,14 +166,14 @@ function StochasticTrainer:train(dataset)
       self.epoch = self.epoch + 1
       currentLearningRate = self.learningRate/(1+self.epoch*self.learningRateDecay)
 
+      if dataset.infiniteSet then
+         self.trainOffset = self.trainOffset + dataset:size()
+      end
+
       if self.maxEpoch > 0 and self.epoch > self.maxEpoch then
          print("<trainer> you have reached the maximum number of epochs")
          break
       end
-
-      if dataset.infiniteSet then
-         self.trainOffset = self.trainOffset + dataset:size()
-      end
    end
 end