Bugfix.

src/lambo/segmenter/lambo.py

@@ -202,6 +202,7 @@ class Lambo():
         with torch.no_grad():
             X = [x.to(self.device) for x in X]
             Y = self.model(*X)
+        Y = Y.to('cpu')
         
         # perform postprocessing
         decisions = self.model.postprocessing(Y, text)