diff --git a/.gitignore b/.gitignore
index e7e232e..5335deb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,3 @@
 __pycache__
 .vscode
-devenv
\ No newline at end of file
+devenv
diff --git a/modules/swap/module.py b/modules/swap/module.py
index 01a0d15..b99e58a 100644
--- a/modules/swap/module.py
+++ b/modules/swap/module.py
@@ -85,7 +85,6 @@ class SwapModule(AbstractModule):
         self._swap_location_check()
         backup_fstab = "sudo cp /etc/fstab /etc/fstab.bak"
         enable_persistence = "echo '{} none swap sw 0 0' | sudo tee -a /etc/fstab".format(self.swap_file)
-
         subprocess.call(backup_fstab, shell=True)
         subprocess.call(enable_persistence, shell=True)
         print("Swap is now persistent!")
@@ -116,9 +115,9 @@ class SwapModule(AbstractModule):
         print("Temporary swapiness is " + str(options[selected_swapiness]))
 
     def _delete_swap(self):
-        location = self._swap_location_check()
-        if location == "NoSwap":
-            return None
+        self._swap_location_check()
+        if self.swap_file is None:
+            return
         else:
             disable_swapfile = "sudo swapoff {} && ".format(self.swap_file) + \
                                "sudo rm {}".format(self.swap_file)