diff --git a/paltas/generate.py b/paltas/generate.py
index 0221fc03..4f5f80c5 100755
--- a/paltas/generate.py
+++ b/paltas/generate.py
@@ -24,7 +24,11 @@
 from paltas.Utils.cli_maker import make_cli
 
 
-def generate_from_config(config_path, save_folder, n=1, save_png_too=False, tf_record=False):
+def generate_from_config(
+		config_path, save_folder, 
+		n: int = 1, 
+		save_png_too: bool = False, 
+		tf_record: bool = False):
 	"""Generate simulated strong lensing images
 
 	Args: