diff --git a/pesto/__init__.py b/pesto/__init__.py
index cff2aebc8285d3f463252869e8f432cb0189bade..5da69bbe1733e6b3bb286735d0dd2a98613ce66a 100644
--- a/pesto/__init__.py
+++ b/pesto/__init__.py
@@ -1 +1,4 @@
 from .core import load_model, predict, predict_from_files
+
+
+__version__ = '1.1.0'
diff --git a/pesto/core.py b/pesto/core.py
index 7c3522510d142d98d7df2310e6f337d468f6dd6f..dd28a5872e8caa8a1f49160743b9d9fe05138a42 100644
--- a/pesto/core.py
+++ b/pesto/core.py
@@ -146,4 +146,3 @@ def predict_from_files(
             predictions = [p.cpu().numpy() for p in predictions]
             for fmt in export_format:
                 export(fmt, output_file, *predictions)
-
diff --git a/pesto/model.py b/pesto/model.py
index a728e3522000b07e07e305d615b6c7f2f66785fd..b9cf5eda18349e38a2d4de1f739798bd5b1df0e6 100644
--- a/pesto/model.py
+++ b/pesto/model.py
@@ -215,10 +215,9 @@ class PESTO(nn.Module):
         if batch_size:
             activations = activations.view(batch_size, -1, activations.size(-1))
 
-        preds = reduce_activations(activations, reduction=self.reduction)
+        activations = activations.roll(-round(self.shift.cpu().item() * self.bins_per_semitone), -1)
 
-        # decrease by shift to get absolute pitch
-        preds.sub_(self.shift)
+        preds = reduce_activations(activations, reduction=self.reduction)
 
         if convert_to_freq:
             preds = 440 * 2 ** ((preds - 69) / 12)
diff --git a/pesto/utils/export.py b/pesto/utils/export.py
index 05c7c5af93008a5e6c57a301ff621c14b740a0c9..70aeb54861b9bff67661fa58793be8e3230ae507 100644
--- a/pesto/utils/export.py
+++ b/pesto/utils/export.py
@@ -44,8 +44,10 @@ def export_png(output_file: str, timesteps, confidence, activations, lims=(21, 1
     activations = activations * confidence[:, None]
     plt.imshow(activations.T,
                aspect='auto', origin='lower', cmap='inferno',
-               extent=(timesteps[0], timesteps[-1]) + lims)
+               extent=(timesteps[0] / 1000, timesteps[-1] / 1000) + lims)
 
+    plt.xlabel("Time (s)")
+    plt.ylabel("Pitch (semitones)")
     plt.title(output_file.rsplit('.', 2)[0])
     plt.tight_layout()
     plt.savefig(output_file)
diff --git a/pyproject.toml b/pyproject.toml
index 6464d322af62df07aeccc2ccf85ecdbde83a636b..b2fb7d05dcc35b499ede3b0562769602baaab66d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pesto-pitch"
-version = "1.0.0"
+dynamic = ["version"]
 authors = [
     {name = "Alain Riou", email = "alain.riou@sony.com"}
 ]
@@ -41,5 +41,8 @@ source = "https://github.com/SonyCSLParis/pesto"
 [tool.pytest.ini_options]
 testpaths = "tests/"
 
+[tool.setuptools.dynamic]
+version = {attr = "pesto.__version__"}
+
 [tool.setuptools.package-data]
 pesto = ["weights/*.ckpt"]