diff --git a/binarycpython/utils/grid.py b/binarycpython/utils/grid.py
index 3d3e71fda9ebd13d1f387f62776359d68fbfc507..b8553a339b87020c50463c25304e92bf8c215deb 100644
--- a/binarycpython/utils/grid.py
+++ b/binarycpython/utils/grid.py
@@ -1637,20 +1637,20 @@ class Population:
                 raise ValueError
 
             # Set up the grid code with a dry run option to see total probability
-            self._generate_grid_code(dry_run=True)
+            if self.grid_options['do_dry_run']:
+                self._generate_grid_code(dry_run=True)
 
-            # Load the grid code
-            self._load_grid_function()
+                # Load the grid code
+                self._load_grid_function()
 
-            if self.grid_options['do_dry_run']:
                 # Do a dry run
                 self._dry_run()
 
-            print(
-                "Total starcount for this run will be: {}".format(
-                    self.grid_options["_total_starcount"]
+                print(
+                    "Total starcount for this run will be: {}".format(
+                        self.grid_options["_total_starcount"]
+                    )
                 )
-            )
 
             #######################
             # Reset values and prepare the grid function
@@ -2308,13 +2308,10 @@ class Population:
 
         if not dry_run:
             # Handling of what is returned, or what is not.
-            # TODO: think of whether this is a good method
-            # code_string += indent * (depth + 2) + "if (self.grid_options['multiplicity'] >= 2): print('phasevol_q: ',phasevol_q); print('phasevol_log10per: ',phasevol_log10per);\n"
             code_string += indent * (depth + 2) + "yield(parameter_dict)\n"
 
         # If its a dry run, dont do anything with it
         else:
-            # code_string += indent * (depth + 2) + "if (self.grid_options['multiplicity'] >= 2): print(phasevol_q)\n"
             code_string += indent * (depth + 2) + "pass\n"
 
         code_string += indent * (depth + 1) + "#" * 40 + "\n"