diff --git a/checks/flash/flake-module.nix b/checks/flash/flake-module.nix
index 2123388b6..2f6725369 100644
--- a/checks/flash/flake-module.nix
+++ b/checks/flash/flake-module.nix
@@ -59,7 +59,7 @@
           name = "flash";
           nodes.target = {
             virtualisation.emptyDiskImages = [ 4096 ];
-            virtualisation.memorySize = 3000;
+            virtualisation.memorySize = 4096;
             environment.systemPackages = [ self.packages.${pkgs.system}.clan-cli ];
             environment.etc."install-closure".source = "${closureInfo}/store-paths";
 
diff --git a/pkgs/clan-cli/clan_lib/machines/update.py b/pkgs/clan-cli/clan_lib/machines/update.py
index 0ec0f8685..b56a8851a 100644
--- a/pkgs/clan-cli/clan_lib/machines/update.py
+++ b/pkgs/clan-cli/clan_lib/machines/update.py
@@ -119,11 +119,9 @@ def run_machine_update(
     with ExitStack() as stack:
         target_host = stack.enter_context(target_host.ssh_control_master())
 
-        if build_host is not None:
+        if build_host:
             build_host = stack.enter_context(build_host.ssh_control_master())
 
-        host = build_host or target_host
-
         sudo_host = stack.enter_context(target_host.become_root())
 
         generate_facts([machine], service=None, regenerate=False)
@@ -132,7 +130,10 @@ def run_machine_update(
         upload_secrets(machine, sudo_host)
         upload_secret_vars(machine, sudo_host)
 
-        path = upload_sources(machine, sudo_host)
+        if build_host:
+            path = upload_sources(machine, build_host)
+        else:
+            path = upload_sources(machine, target_host)
 
         nix_options = machine.flake.nix_options if machine.flake.nix_options else []
 
@@ -174,11 +175,13 @@ def run_machine_update(
                 *nix_options,
             ]
 
-        if become_root:
-            host = sudo_host
+        if become_root and not build_host:
+            target_host = sudo_host
 
-        remote_env = host.nix_ssh_env(control_master=False)
-        ret = host.run(
+        deploy_host = build_host if build_host else target_host
+
+        remote_env = deploy_host.nix_ssh_env(control_master=False)
+        ret = deploy_host.run(
             switch_cmd,
             RunOpts(
                 check=False,
@@ -205,7 +208,7 @@ def run_machine_update(
                 machine.info(
                     "Mobile machine detected, applying workaround deployment method"
                 )
-            ret = host.run(
+            ret = deploy_host.run(
                 ["nixos--rebuild", "test", *nix_options] if is_mobile else switch_cmd,
                 RunOpts(
                     log=Log.BOTH,