From c76c3e52b3c8a9c793042df38ef0d144f2229a30 Mon Sep 17 00:00:00 2001
From: piotrmp <piotr.m.przybyla@gmail.com>
Date: Tue, 17 Jan 2023 14:53:38 +0100
Subject: [PATCH] Bugfix.

---
 src/lambo/segmenter/lambo.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/lambo/segmenter/lambo.py b/src/lambo/segmenter/lambo.py
index 134f28f..e481fb9 100644
--- a/src/lambo/segmenter/lambo.py
+++ b/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)
-- 
GitLab