DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Removed for ZeroGPU compatibility (lazy check instead)

main.py

@@ -33,7 +33,7 @@ except ImportError:
 # ----------------------------------------------------------------------
 # CONFIGURATION
 # ----------------------------------------------------------------------
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"  # Removed for ZeroGPU compatibility (lazy check instead)
 
 # Model options
 MODEL_SIZE = 1024
@@ -1252,6 +1252,9 @@ if __name__ == "__main__":
         num_workers = max(1, min(NUM_WORKERS, total_cpus))
     
     # Decide whether to use multiprocessing
+    # Local device check
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    
     use_pool = USE_MULTIPROCESSING and DEVICE == "cpu" and total_cpus >= 4
     
     if use_pool: