TF2.0 have some problem with generator length > steps_per_epoch: tens…

…orflow/tensorflow#25254
notha99y · Jun 15, 2020 · c7519aa · c7519aa
1 parent 12f929f
commit c7519aa
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/train.py b/train.py
@@ -67,7 +67,7 @@ def simple_cnn():
 
 
 def get_resnet50():
-    height, width, channel = (240, 320, 13)
+    height, width, channel = (64, 64, 13)
     num_of_classes = 10
 
     base_model = ResNet50(
@@ -120,8 +120,8 @@ def get_resnet50():
         shuffle=True,
     )
 
-    model.fit(
+    model.fit_generator(
         train_gen,
         epochs=25,
-        steps_per_epoch=train_gen.n // train_gen.batch_size + 1,
+        steps_per_epoch=train_gen.n // train_gen.batch_size,
     )