Update app.py

app.py

@@ -29,6 +29,8 @@ def inference(point1, point2, point3, point4):
     w = torch.stack(w)
     if len(pts.shape) == 2:
         pts = pts[None]
+    if len(w.shape) == 1:
+        w = w[None]
     
     l2 = fit_line(pts, w)