diff --git a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.model.py.swp b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.model.py.swp
index 9c52f6176..904af1781 100644
Binary files a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.model.py.swp and b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.model.py.swp differ
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/model.py b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/model.py
index a6e919e2f..9f2305393 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/model.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/model.py
@@ -113,7 +113,7 @@ class Transducer(nn.Module):
               lm_scale * lm_probs + am_scale * am_probs +
               (1-lm_scale-am_scale) * combined_probs
         """
-        assert x.ndim == 3, x.shape
+        assert x.ndim == 2 or x.ndim == 3, x.shape
         assert x_lens.ndim == 1, x_lens.shape
         assert y.num_axes == 2, y.num_axes