Update app.py

app.py

@@ -10,6 +10,7 @@ DEVICE = torch.device("cpu")
 
 # ---------------- LOAD ENCODER ----------------
 encoder = load_encoder("encoder_resnet18_simclr.pth")
+encoder.eval()
 
 # ---------------- LOAD LINEAR PROBE ----------------
 data = np.load("linear_probe_cifar10.npz", allow_pickle=True)
@@ -33,12 +34,15 @@ transform = transforms.Compose([
 # ---------------- PREDICT ----------------
 def predict(image):
     image = Image.fromarray(image).convert("RGB")
-    x = transform(image).unsqueeze(0)
+    x = transform(image).unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
-        emb = encoder(x).numpy()
+        emb = encoder(x).cpu().numpy()
 
-    emb = torch.tensor(emb)
+    # IMPORTANT — match training behavior
+    emb = emb / np.linalg.norm(emb, axis=1, keepdims=True)
+
+    emb = torch.tensor(emb, dtype=torch.float32)
 
     logits = emb @ W.T + b
     probs = F.softmax(logits, dim=1).numpy()[0]
@@ -54,4 +58,5 @@ demo = gr.Interface(
     description="Upload an image. Model trained WITHOUT labels using SimCLR. Evaluated using Linear Probe."
 )
 
-demo.launch()
+demo.queue(concurrency_count=4)
+demo.launch(server_name="0.0.0.0", server_port=7860)