Update ball.py

ball.py

@@ -44,8 +44,8 @@ import datetime
 class TrainingConfig:
     image_size = 256  # the generated image resolution
     train_batch_size = 1
-    eval_batch_size = 100 # how many images to sample during evaluation
-    num_epochs = 2
+    eval_batch_size = 50 # how many images to sample during evaluation
+    num_epochs = 1
     gradient_accumulation_steps = 1
     learning_rate = 1e-4 #Was 1e-4
     lr_warmup_steps = 500
@@ -228,7 +228,7 @@ def evaluate(config, epoch, pipeline):
     ).images
 
     # Make a grid out of the images
-    image_grid = make_grid(images, rows=10, cols=10)
+    image_grid = make_grid(images, rows=10, cols=5)
 
     # Save the images
     test_dir = os.path.join(config.output_dir, "samples"+config.time_started)