diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py b/egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py
index 9bac46004..bcd419fb7 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py
@@ -849,6 +849,8 @@ class RelPositionalEncoding(torch.nn.Module):
             torch.Tensor: Encoded tensor (batch, 2*time-1, `*`).
 
         """
+        if isinstance(left_context, torch.Tensor):
+            left_context = left_context.item()
         self.extend_pe(x, left_context)
         x_size_1 = x.size(1) + left_context
         pos_emb = self.pe[
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless3/test_onnx.py b/egs/librispeech/ASR/pruned_transducer_stateless3/test_onnx.py
index 598fcf344..810da8da6 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless3/test_onnx.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless3/test_onnx.py
@@ -113,7 +113,7 @@ def test_rel_pos():
 
     torch.onnx.export(
         encoder_pos,
-        x,
+        (x, torch.zeros(1, dtype=torch.int64)),
         filename,
         verbose=False,
         opset_version=opset_version,
@@ -139,7 +139,9 @@ def test_rel_pos():
     assert input_nodes[0].name == "x"
     assert input_nodes[0].shape == ["N", "T", num_features]
 
-    inputs = {input_nodes[0].name: x.numpy()}
+    inputs = {
+        input_nodes[0].name: x.numpy(),
+    }
     onnx_y, onnx_pos_emb = session.run(["y", "pos_emb"], inputs)
     onnx_y = torch.from_numpy(onnx_y)
     onnx_pos_emb = torch.from_numpy(onnx_pos_emb)