diff --git a/penn/decode.py b/penn/decode.py
index 2722293..2cd4fb1 100644
--- a/penn/decode.py
+++ b/penn/decode.py
@@ -73,11 +73,11 @@ def __call__(self, logits):
         gpu = (
             None if distributions.device.type == 'cpu'
             else distributions.device.index)
-        bins = torbi.decode(
-            distributions[0].T,
-            self.transition,
-            self.initial,
-            gpu)
+        bins = torbi.from_probabilities(
+            observation=distributions[0].T.unsqueeze(dim=0),
+            transition=self.transition,
+            initial=self.initial,
+            gpu=gpu)
 
         # Convert to frequency in Hz
         if self.local_expected_value:
@@ -131,11 +131,11 @@ def __call__(self, logits):
         gpu = (
             None if distributions.device.type == 'cpu'
             else distributions.device.index)
-        bins = torbi.decode(
-            distributions[0].T,
-            self.transition,
-            self.initial,
-            gpu)
+        bins = torbi.from_probabilities(
+            observation=distributions[0].T.unsqueeze(dim=0),
+            transition=self.transition,
+            initial=self.initial,
+            gpu=gpu)
 
         # Convert to frequency in Hz
         if self.local_expected_value: