Update app.py

app.py

@@ -27,7 +27,9 @@ def inference(point1, point2, point3, point4):
         w.append(1 - p2_noise.mean())
     pts = torch.stack(pts)
     w = torch.stack(w)
-
+    if len(pts.shape) == 2:
+        pts = pts[None]
+    
     l2 = fit_line(pts, w)
     
     # project some points along the estimated line