diff --git a/rhythmblock.lua b/rhythmblock.lua
index 1899de3..13543c8 100644
--- a/rhythmblock.lua
+++ b/rhythmblock.lua
@@ -4,13 +4,17 @@ buttonColor = colors.lightGray
 textColor = colors.lightGray
 altTextColor = colors.gray
 backgroundColor = colors.black
-if arg[1] == nil then
+
+-- Code starts from here
+
 function round(n) -- We need a function to round the center of the terminal
     return n % 1 >= 0.5 and math.ceil(n) or math.floor(n)
 end
 width, height = term.getSize() -- Gets the terminal size to determine the center
 centerWidth = round(width / 2) -- Defines the horizontal center
 centerHeight = round(height / 2) -- Defines the vertical center
+
+if arg[1] == nil then
 	peripherals = peripheral.getNames() -- Gets peripherals to check if any disk drives are avavailable
 	if #peripherals < 0 then
 		print "No drive"
@@ -26,7 +30,7 @@ centerHeight = round(height / 2) -- Defines the vertical center
 		end
 		if driveCount > 1 then
 			print("Too many disk drives. Specify where the disk drive is by running rhythmbox [drive position]") -- For safety reasons
-			os.exit()
+			return false	
 		end
 	end
 else