diff --git a/WoeUSB/core.py b/WoeUSB/core.py
index 28dbb49..cb2057a 100644
--- a/WoeUSB/core.py
+++ b/WoeUSB/core.py
@@ -108,7 +108,7 @@ def init(from_cli=True, install_mode=None, source_media=None, target_media=None,
         return [source_fs_mountpoint, target_fs_mountpoint, temp_directory, install_mode, source_media, target_media,
                 workaround_bios_boot_flag, target_filesystem_type, filesystem_label, verbose, debug, parser]
     else:
-        return [source_fs_mountpoint, target_fs_mountpoint, temp_directory]
+        return [source_fs_mountpoint, target_fs_mountpoint, temp_directory, target_media]
 
 
 def main(source_fs_mountpoint, target_fs_mountpoint, source_media, target_media, install_mode, temp_directory,
diff --git a/WoeUSB/gui.py b/WoeUSB/gui.py
index 00e03b5..0a3b7a7 100644
--- a/WoeUSB/gui.py
+++ b/WoeUSB/gui.py
@@ -393,7 +393,7 @@ def __init__(self, source, target, boot_flag, filesystem):
         self.filesystem = filesystem
 
     def run(self):
-        source_fs_mountpoint, target_fs_mountpoint, temp_directory = core.init(
+        source_fs_mountpoint, target_fs_mountpoint, temp_directory, target_media = core.init(
             from_cli=False,
             install_mode="device",
             source_media=self.source,
@@ -405,7 +405,7 @@ def run(self):
         except SystemExit:
             pass
 
-        core.cleanup(source_fs_mountpoint, target_fs_mountpoint, temp_directory)
+        core.cleanup(source_fs_mountpoint, target_fs_mountpoint, temp_directory, target_media)
 
 
 def run():