diff --git a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.train.py.swp b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.train.py.swp
new file mode 100644
index 000000000..7676e6c3d
Binary files /dev/null and b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.train.py.swp differ
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py
index c211427e7..376fa999d 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py
@@ -672,7 +672,7 @@ def compute_loss(
     device = model.device if isinstance(model, DDP) else next(model.parameters()).device
     feature = batch["inputs"]
     # at entry, feature is (N, T, C)
-    assert feature.ndim == 3
+    assert feature.ndim == 2 or feature.ndim == 3
     feature = feature.to(device)
 
     supervisions = batch["supervisions"]