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