Fix offload on interrupt

nodes_sampler.py

@@ -2593,10 +2593,10 @@ class WanVideoSampler:
 
             except Exception as e:
                 log.error(f"Error during sampling: {e}")
-                if force_offload:
-                    if not model["auto_cpu_offload"]:
+                raise
+            finally:
+                if force_offload and not model["auto_cpu_offload"]:
                     offload_transformer(transformer)
-                raise e
 
         if phantom_latents is not None:
             latent = latent[:,:-phantom_latents.shape[1]]
@@ -2620,10 +2620,6 @@ class WanVideoSampler:
                     "magcache_state": transformer.magcache_state,
                 }
 
-        if force_offload:
-            if not model["auto_cpu_offload"]:
-                offload_transformer(transformer)
-
         try:
             print_memory(device)
             torch.cuda.reset_peak_memory_stats(device)