Change version of loading model

app.py

@@ -27,20 +27,18 @@ with open("token_probabilities.json") as f:
 token_probabilities = np.array([token_probs_dict[str(i)] for i in range(len(token_probs_dict))], dtype=np.float32)
 
 @spaces.GPU
-def load_model():
+def load_weights():
+    # OK: download & load weights to CPU
     ckpt_path = hf_hub_download(
         repo_id="ruurd/tini_model",
         filename="diffusion-model.pth",
         token=os.getenv("HF_TOKEN")
     )
+    return torch.load(ckpt_path, map_location="cpu")  # ✅ returns only CPU tensors
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = torch.load(ckpt_path, map_location=device)
-    model = disable_dropout(model)
-    model.to(device)
-    model.eval()
-    return model
-
+model = CustomTransformerModel(...)
+model.load_state_dict(load_weights())
+model.to("cuda")  # ✅ OK now, after @spaces.GPU is done
 
 rng = np.random.default_rng()
 
@@ -146,10 +144,6 @@ def diffusion_chat(question, eot_weight, max_it, sharpness):
 
 # --- Gradio Interface ---
 
-print("Loading model...")
-model = load_model()
-print("✅ Model loaded.")
-
 demo = gr.Interface(
     fn=diffusion_chat,
     inputs=[