diff --git a/MultiSource/Applications/viterbi/CMakeLists.txt b/MultiSource/Applications/viterbi/CMakeLists.txt
--- a/MultiSource/Applications/viterbi/CMakeLists.txt
+++ b/MultiSource/Applications/viterbi/CMakeLists.txt
@@ -1,3 +1,8 @@
 list(APPEND CPPFLAGS -I${CMAKE_CURRENT_SOURCE_DIR} -DPROJ_SRC_DIR=".")
+if(SMALL_PROBLEM_SIZE)
+  set(INPUTS Dist_demux_small)
+else()
+  set(INPUTS Dist_demux)
+endif()
 llvm_multisource(viterbi)
-llvm_test_data(viterbi Dist_demux)
+llvm_test_data(viterbi ${INPUTS})