Update butterfly_training.py

butterfly_training.py

@@ -61,12 +61,12 @@ class TrainingConfig:
     image_size = 64  # the generated image resolution
     train_batch_size = 16
     eval_batch_size = 16  # how many images to sample during evaluation
-    num_epochs = 50
+    num_epochs = 10
     gradient_accumulation_steps = 1
     learning_rate = 1e-4
     lr_warmup_steps = 500
-    save_image_epochs = 10
-    save_model_epochs = 30
+    save_image_epochs = 5
+    save_model_epochs = 5
     mixed_precision = "fp16"  # `no` for float32, `fp16` for automatic mixed precision
     output_dir = "butterflystest"  # the model name locally and on the HF Hub